// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
// SPDX-License-Identifier: MIT

package handshake

import (
	"testing"

	"github.com/pion/dtls/v3/internal/ciphersuite/types"
	"github.com/pion/dtls/v3/pkg/crypto/elliptic"
	"github.com/pion/dtls/v3/pkg/crypto/hash"
	"github.com/pion/dtls/v3/pkg/crypto/signature"
	"github.com/stretchr/testify/assert"
)

func TestHandshakeMessageServerKeyExchange(t *testing.T) {
	test := func(rawServerKeyExchange []byte, parsedServerKeyExchange *MessageServerKeyExchange) {
		c := &MessageServerKeyExchange{
			KeyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
		}
		assert.NoError(t, c.Unmarshal(rawServerKeyExchange))
		assert.Equal(t, parsedServerKeyExchange, c)

		raw, err := c.Marshal()
		assert.NoError(t, err)
		assert.Equal(t, rawServerKeyExchange, raw)
	}

	t.Run("Hash+Signature", func(*testing.T) {
		rawServerKeyExchange := []byte{
			0x03, 0x00, 0x1d, 0x41, 0x04, 0x0c, 0xb9, 0xa3, 0xb9, 0x90, 0x71, 0x35, 0x4a, 0x08, 0x66, 0xaf,
			0xd6, 0x88, 0x58, 0x29, 0x69, 0x98, 0xf1, 0x87, 0x0f, 0xb5, 0xa8, 0xcd, 0x92, 0xf6, 0x2b, 0x08,
			0x0c, 0xd4, 0x16, 0x5b, 0xcc, 0x81, 0xf2, 0x58, 0x91, 0x8e, 0x62, 0xdf, 0xc1, 0xec, 0x72, 0xe8,
			0x47, 0x24, 0x42, 0x96, 0xb8, 0x7b, 0xee, 0xe7, 0x0d, 0xdc, 0x44, 0xec, 0xf3, 0x97, 0x6b, 0x1b,
			0x45, 0x28, 0xac, 0x3f, 0x35, 0x02, 0x03, 0x00, 0x47, 0x30, 0x45, 0x02, 0x21, 0x00, 0xb2, 0x0b,
			0x22, 0x95, 0x3d, 0x56, 0x57, 0x6a, 0x3f, 0x85, 0x30, 0x6f, 0x55, 0xc3, 0xf4, 0x24, 0x1b, 0x21,
			0x07, 0xe5, 0xdf, 0xba, 0x24, 0x02, 0x68, 0x95, 0x1f, 0x6e, 0x13, 0xbd, 0x9f, 0xaa, 0x02, 0x20,
			0x49, 0x9c, 0x9d, 0xdf, 0x84, 0x60, 0x33, 0x27, 0x96, 0x9e, 0x58, 0x6d, 0x72, 0x13, 0xe7, 0x3a,
			0xe8, 0xdf, 0x43, 0x75, 0xc7, 0xb9, 0x37, 0x6e, 0x90, 0xe5, 0x3b, 0x81, 0xd4, 0xda, 0x68, 0xcd,
		}
		parsedServerKeyExchange := &MessageServerKeyExchange{
			EllipticCurveType:    elliptic.CurveTypeNamedCurve,
			NamedCurve:           elliptic.X25519,
			PublicKey:            rawServerKeyExchange[4:69],
			HashAlgorithm:        hash.SHA1,
			SignatureAlgorithm:   signature.ECDSA,
			Signature:            rawServerKeyExchange[73:144],
			KeyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
		}

		test(rawServerKeyExchange, parsedServerKeyExchange)
	})

	t.Run("Anonymous", func(*testing.T) {
		rawServerKeyExchange := []byte{
			0x03, 0x00, 0x1d, 0x41, 0x04, 0x0c, 0xb9, 0xa3, 0xb9, 0x90, 0x71, 0x35, 0x4a, 0x08, 0x66, 0xaf,
			0xd6, 0x88, 0x58, 0x29, 0x69, 0x98, 0xf1, 0x87, 0x0f, 0xb5, 0xa8, 0xcd, 0x92, 0xf6, 0x2b, 0x08,
			0x0c, 0xd4, 0x16, 0x5b, 0xcc, 0x81, 0xf2, 0x58, 0x91, 0x8e, 0x62, 0xdf, 0xc1, 0xec, 0x72, 0xe8,
			0x47, 0x24, 0x42, 0x96, 0xb8, 0x7b, 0xee, 0xe7, 0x0d, 0xdc, 0x44, 0xec, 0xf3, 0x97, 0x6b, 0x1b,
			0x45, 0x28, 0xac, 0x3f, 0x35,
		}
		parsedServerKeyExchange := &MessageServerKeyExchange{
			EllipticCurveType:    elliptic.CurveTypeNamedCurve,
			NamedCurve:           elliptic.X25519,
			PublicKey:            rawServerKeyExchange[4:69],
			HashAlgorithm:        hash.None,
			SignatureAlgorithm:   signature.Anonymous,
			KeyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
		}

		test(rawServerKeyExchange, parsedServerKeyExchange)
	})
}

func TestHandshakeMessageServerKeyExchangeUnmarshalErrors(t *testing.T) {
	for _, test := range []struct {
		name                 string
		keyExchangeAlgorithm types.KeyExchangeAlgorithm
		data                 []byte
		expectedErr          error
	}{
		{
			name:                 "BufferTooSmall",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x00},
			expectedErr:          errBufferTooSmall,
		},
		{
			name:                 "CipherSuiteUnset",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmNone,
			data:                 []byte{0x00, 0x00},
			expectedErr:          errCipherSuiteUnset,
		},
		{
			// PSK-only: a non-empty body remains after the identity hint.
			name:                 "PskLengthMismatch",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmPsk,
			data:                 []byte{0x00, 0x00, 0xAA},
			expectedErr:          errLengthMismatch,
		},
		{
			// An algorithm that is neither PSK nor ECDHE is unsupported here.
			name:                 "UnsupportedKeyExchangeAlgorithm",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithm(1),
			data:                 []byte{0x00, 0x00},
			expectedErr:          errLengthMismatch,
		},
		{
			// ECDHE_PSK where the (empty) identity hint consumes the whole body,
			// leaving nothing for the ECDHE parameters. Previously panicked.
			name:                 "EcdhePskEmptyHintConsumesBody",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmPsk | types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x00, 0x00},
			expectedErr:          errBufferTooSmall,
		},
		{
			name:                 "InvalidEllipticCurveType",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x99, 0x00},
			expectedErr:          errInvalidEllipticCurveType,
		},
		{
			// Valid curve type but not enough bytes for the named curve.
			name:                 "NamedCurveBufferTooSmall",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00},
			expectedErr:          errBufferTooSmall,
		},
		{
			name:                 "InvalidNamedCurve",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0xFF, 0xFF},
			expectedErr:          errInvalidNamedCurve,
		},
		{
			// Valid curve type and named curve but missing the public key length.
			name:                 "PublicKeyLengthBufferTooSmall",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00, 0x1d},
			expectedErr:          errBufferTooSmall,
		},
		{
			// Public key length exceeds the remaining body.
			name:                 "PublicKeyBufferTooSmall",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00, 0x1d, 0x05},
			expectedErr:          errBufferTooSmall,
		},
		{
			// Valid signature algorithm (0x03 = ECDSA) but invalid hash algorithm (0x07).
			name:                 "InvalidHashAlgorithm",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00, 0x1d, 0x00, 0x07, 0x03},
			expectedErr:          errInvalidSignHashAlgorithm,
		},
		{
			// Valid hash algorithm but missing the signature algorithm byte.
			name:                 "SignatureAlgorithmBufferTooSmall",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00, 0x1d, 0x00, 0x04},
			expectedErr:          errBufferTooSmall,
		},
		{
			// Valid hash algorithm (0x04 = SHA256) but invalid signature algorithm (0x02).
			name:                 "InvalidSignatureAlgorithm",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00, 0x1d, 0x00, 0x04, 0x02},
			expectedErr:          errInvalidSignHashAlgorithm,
		},
		{
			// Valid hash and signature algorithms but missing the signature length.
			name:                 "SignatureLengthBufferTooSmall",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00, 0x1d, 0x00, 0x04, 0x03},
			expectedErr:          errBufferTooSmall,
		},
		{
			// Signature length exceeds the remaining body.
			name:                 "SignatureBufferTooSmall",
			keyExchangeAlgorithm: types.KeyExchangeAlgorithmEcdhe,
			data:                 []byte{0x03, 0x00, 0x1d, 0x00, 0x04, 0x03, 0x00, 0x05},
			expectedErr:          errBufferTooSmall,
		},
	} {
		c := &MessageServerKeyExchange{
			KeyExchangeAlgorithm: test.keyExchangeAlgorithm,
		}
		assert.ErrorIs(t, c.Unmarshal(test.data), test.expectedErr, test.name)
	}
}
