package qtls

import (
	"bytes"
	"fmt"
	"net"
	"testing"
	"time"
)

type exportedKey struct {
	typ           string // "read" or "write"
	encLevel      EncryptionLevel
	suite         *CipherSuiteTLS13
	trafficSecret []byte
}

func compareExportedKeys(t *testing.T, k1, k2 *exportedKey) {
	if k1.encLevel != k2.encLevel || k1.suite.ID != k2.suite.ID || !bytes.Equal(k1.trafficSecret, k2.trafficSecret) {
		t.Fatal("mismatching keys")
	}
}

type recordLayerWithKeys struct {
	in  <-chan []byte
	out chan<- interface{}
}

func (r *recordLayerWithKeys) SetReadKey(encLevel EncryptionLevel, suite *CipherSuiteTLS13, trafficSecret []byte) {
	r.out <- &exportedKey{typ: "read", encLevel: encLevel, suite: suite, trafficSecret: trafficSecret}
}
func (r *recordLayerWithKeys) SetWriteKey(encLevel EncryptionLevel, suite *CipherSuiteTLS13, trafficSecret []byte) {
	r.out <- &exportedKey{typ: "write", encLevel: encLevel, suite: suite, trafficSecret: trafficSecret}
}
func (r *recordLayerWithKeys) ReadHandshakeMessage() ([]byte, error) { return <-r.in, nil }
func (r *recordLayerWithKeys) WriteRecord(b []byte) (int, error)     { r.out <- b; return len(b), nil }
func (r *recordLayerWithKeys) SendAlert(uint8)                       {}

type unusedConn struct {
	remoteAddr net.Addr
}

var _ net.Conn = &unusedConn{}

func (unusedConn) Read([]byte) (int, error)         { panic("unexpected call to Read()") }
func (unusedConn) Write([]byte) (int, error)        { panic("unexpected call to Write()") }
func (unusedConn) Close() error                     { return nil }
func (unusedConn) LocalAddr() net.Addr              { return &net.TCPAddr{} }
func (c *unusedConn) RemoteAddr() net.Addr          { return c.remoteAddr }
func (unusedConn) SetDeadline(time.Time) error      { return nil }
func (unusedConn) SetReadDeadline(time.Time) error  { return nil }
func (unusedConn) SetWriteDeadline(time.Time) error { return nil }

func TestAlternativeRecordLayer(t *testing.T) {
	sIn := make(chan []byte, 10)
	sOut := make(chan interface{}, 10)
	defer close(sOut)
	cIn := make(chan []byte, 10)
	cOut := make(chan interface{}, 10)
	defer close(cOut)

	serverEvents := make(chan interface{}, 100)
	go func() {
		for {
			c, ok := <-sOut
			if !ok {
				return
			}
			serverEvents <- c
			if b, ok := c.([]byte); ok {
				cIn <- b
			}
		}
	}()

	clientEvents := make(chan interface{}, 100)
	go func() {
		for {
			c, ok := <-cOut
			if !ok {
				return
			}
			clientEvents <- c
			if b, ok := c.([]byte); ok {
				sIn <- b
			}
		}
	}()

	errChan := make(chan error)
	go func() {
		config := testConfig.Clone()
		config.AlternativeRecordLayer = &recordLayerWithKeys{in: sIn, out: sOut}
		tlsConn := Server(&unusedConn{}, config)
		defer tlsConn.Close()
		errChan <- tlsConn.Handshake()
	}()

	config := testConfig.Clone()
	config.AlternativeRecordLayer = &recordLayerWithKeys{in: cIn, out: cOut}
	tlsConn := Client(&unusedConn{}, config)
	defer tlsConn.Close()
	if err := tlsConn.Handshake(); err != nil {
		t.Fatalf("Handshake failed: %s", err)
	}

	// Handshakes completed. Now check that events were received in the correct order.
	var clientHandshakeReadKey, clientHandshakeWriteKey *exportedKey
	var clientApplicationReadKey, clientApplicationWriteKey *exportedKey
	for i := 0; i <= 5; i++ {
		ev := <-clientEvents
		switch i {
		case 0:
			if ev.([]byte)[0] != typeClientHello {
				t.Fatalf("expected ClientHello")
			}
		case 1:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "write" || keyEv.encLevel != EncryptionHandshake {
				t.Fatalf("expected the handshake write key")
			}
			clientHandshakeWriteKey = keyEv
		case 2:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "read" || keyEv.encLevel != EncryptionHandshake {
				t.Fatalf("expected the handshake read key")
			}
			clientHandshakeReadKey = keyEv
		case 3:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "read" || keyEv.encLevel != EncryptionApplication {
				t.Fatalf("expected the application read key")
			}
			clientApplicationReadKey = keyEv
		case 4:
			if ev.([]byte)[0] != typeFinished {
				t.Fatalf("expected Finished")
			}
		case 5:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "write" || keyEv.encLevel != EncryptionApplication {
				t.Fatalf("expected the application write key")
			}
			clientApplicationWriteKey = keyEv
		}
	}
	if len(clientEvents) > 0 {
		t.Fatal("didn't expect any more client events")
	}

	for i := 0; i <= 8; i++ {
		ev := <-serverEvents
		switch i {
		case 0:
			if ev.([]byte)[0] != typeServerHello {
				t.Fatalf("expected ServerHello")
			}
		case 1:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "read" || keyEv.encLevel != EncryptionHandshake {
				t.Fatalf("expected the handshake read key")
			}
			compareExportedKeys(t, clientHandshakeWriteKey, keyEv)
		case 2:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "write" || keyEv.encLevel != EncryptionHandshake {
				t.Fatalf("expected the handshake write key")
			}
			compareExportedKeys(t, clientHandshakeReadKey, keyEv)
		case 3:
			if ev.([]byte)[0] != typeEncryptedExtensions {
				t.Fatalf("expected EncryptedExtensions")
			}
		case 4:
			if ev.([]byte)[0] != typeCertificate {
				t.Fatalf("expected Certificate")
			}
		case 5:
			if ev.([]byte)[0] != typeCertificateVerify {
				t.Fatalf("expected CertificateVerify")
			}
		case 6:
			if ev.([]byte)[0] != typeFinished {
				t.Fatalf("expected Finished")
			}
		case 7:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "write" || keyEv.encLevel != EncryptionApplication {
				t.Fatalf("expected the application write key")
			}
			compareExportedKeys(t, clientApplicationReadKey, keyEv)
		case 8:
			keyEv := ev.(*exportedKey)
			if keyEv.typ != "read" || keyEv.encLevel != EncryptionApplication {
				t.Fatalf("expected the application read key")
			}
			compareExportedKeys(t, clientApplicationWriteKey, keyEv)
		}
	}
	if len(serverEvents) > 0 {
		t.Fatal("didn't expect any more server events")
	}
}

func TestForbiddenZeroRTT(t *testing.T) {
	// run the first handshake to get a session ticket
	clientConn, serverConn := localPipe(t)
	errChan := make(chan error, 1)
	go func() {
		tlsConn := Server(serverConn, testConfig.Clone())
		defer tlsConn.Close()
		err := tlsConn.Handshake()
		errChan <- err
		if err != nil {
			return
		}
		tlsConn.Write([]byte{0})
	}()

	clientConfig := testConfig.Clone()
	clientConfig.ClientSessionCache = NewLRUClientSessionCache(10)
	tlsConn := Client(clientConn, clientConfig)
	if err := tlsConn.Handshake(); err != nil {
		t.Fatalf("first handshake failed: %s", err)
	}
	tlsConn.Read([]byte{0}) // make sure to read the session ticket
	tlsConn.Close()
	if err := <-errChan; err != nil {
		t.Fatalf("first handshake failed: %s", err)
	}

	sIn := make(chan []byte, 10)
	sOut := make(chan interface{}, 10)
	defer close(sOut)
	cIn := make(chan []byte, 10)
	cOut := make(chan interface{}, 10)
	defer close(cOut)

	serverEvents := make(chan interface{}, 100)
	go func() {
		for {
			c, ok := <-sOut
			if !ok {
				return
			}
			serverEvents <- c
			if b, ok := c.([]byte); ok {
				cIn <- b
			}
		}
	}()

	clientEvents := make(chan interface{}, 100)
	go func() {
		for {
			c, ok := <-cOut
			if !ok {
				return
			}
			clientEvents <- c
			if b, ok := c.([]byte); ok {
				if b[0] == typeClientHello {
					msg := &clientHelloMsg{}
					if ok := msg.unmarshal(b); !ok {
						panic("unmarshaling failed")
					}
					msg.earlyData = true
					msg.raw = nil
					b = msg.marshal()
				}
				sIn <- b
			}
		}
	}()

	done := make(chan struct{})
	go func() {
		defer close(done)
		clientConfig.AlternativeRecordLayer = &recordLayerWithKeys{in: cIn, out: cOut}
		Client(&unusedConn{remoteAddr: clientConn.RemoteAddr()}, clientConfig).Handshake()
	}()

	config := testConfig.Clone()
	config.MinVersion = VersionTLS13
	config.AlternativeRecordLayer = &recordLayerWithKeys{in: sIn, out: sOut}
	tlsConn = Server(&unusedConn{}, config)
	err := tlsConn.Handshake()
	if err == nil {
		t.Fatal("expected handshake to fail")
	}
	if err.Error() != "tls: client sent unexpected early data" {
		t.Fatalf("expected early data error")
	}
	cIn <- []byte{0} // make the client handshake error
	<-done
}

func TestZeroRTTKeys(t *testing.T) {
	// run the first handshake to get a session ticket
	clientConn, serverConn := localPipe(t)
	errChan := make(chan error, 1)
	go func() {
		config := testConfig.Clone()
		config.MaxEarlyData = 1000
		tlsConn := Server(serverConn, config)
		defer tlsConn.Close()
		err := tlsConn.Handshake()
		errChan <- err
		if err != nil {
			return
		}
		tlsConn.Write([]byte{0})
	}()

	clientConfig := testConfig.Clone()
	clientConfig.ClientSessionCache = NewLRUClientSessionCache(10)
	tlsConn := Client(clientConn, clientConfig)
	if err := tlsConn.Handshake(); err != nil {
		t.Fatalf("first handshake failed: %s", err)
	}
	tlsConn.Read([]byte{0}) // make sure to read the session ticket
	tlsConn.Close()
	if err := <-errChan; err != nil {
		t.Fatalf("first handshake failed: %s", err)
	}

	sIn := make(chan []byte, 10)
	sOut := make(chan interface{}, 10)
	defer close(sOut)
	cIn := make(chan []byte, 10)
	cOut := make(chan interface{}, 10)
	defer close(cOut)

	var serverEarlyData bool
	var serverExportedKey *exportedKey
	go func() {
		for {
			c, ok := <-sOut
			if !ok {
				return
			}
			if b, ok := c.([]byte); ok {
				if b[0] == typeEncryptedExtensions {
					var msg encryptedExtensionsMsg
					if ok := msg.unmarshal(b); !ok {
						panic("failed to unmarshal EncryptedExtensions")
					}
					serverEarlyData = msg.earlyData
				}
				cIn <- b
			}
			if k, ok := c.(*exportedKey); ok && k.encLevel == Encryption0RTT {
				serverExportedKey = k
			}
		}
	}()

	var clientEarlyData bool
	var clientExportedKey *exportedKey
	go func() {
		for {
			c, ok := <-cOut
			if !ok {
				return
			}
			if b, ok := c.([]byte); ok {
				if b[0] == typeClientHello {
					var msg clientHelloMsg
					if ok := msg.unmarshal(b); !ok {
						panic("failed to unmarshal ClientHello")
					}
					clientEarlyData = msg.earlyData
				}
				sIn <- b
			}
			if k, ok := c.(*exportedKey); ok && k.encLevel == Encryption0RTT {
				clientExportedKey = k
			}
		}
	}()

	errChan = make(chan error)
	go func() {
		config := testConfig.Clone()
		config.AlternativeRecordLayer = &recordLayerWithKeys{in: sIn, out: sOut}
		config.MaxEarlyData = 1
		config.Accept0RTT = func([]byte) bool { return true }
		tlsConn := Server(&unusedConn{}, config)
		defer tlsConn.Close()
		errChan <- tlsConn.Handshake()
	}()

	clientConfig.AlternativeRecordLayer = &recordLayerWithKeys{in: cIn, out: cOut}
	clientConfig.Enable0RTT = true
	tlsConn = Client(&unusedConn{remoteAddr: clientConn.RemoteAddr()}, clientConfig)
	defer tlsConn.Close()
	if err := tlsConn.Handshake(); err != nil {
		t.Fatalf("Handshake failed: %s", err)
	}
	if err := <-errChan; err != nil {
		t.Fatalf("Handshake failed: %s", err)
	}

	if !clientEarlyData {
		t.Fatal("expected the client to offer early data")
	}
	if !serverEarlyData {
		t.Fatal("expected the server to offer early data")
	}
	compareExportedKeys(t, clientExportedKey, serverExportedKey)
}

type recordLayer struct {
	in  <-chan []byte
	out chan<- []byte
}

func (r *recordLayer) SetReadKey(encLevel EncryptionLevel, suite *CipherSuiteTLS13, trafficSecret []byte) {
}
func (r *recordLayer) SetWriteKey(encLevel EncryptionLevel, suite *CipherSuiteTLS13, trafficSecret []byte) {
}
func (r *recordLayer) ReadHandshakeMessage() ([]byte, error) { return <-r.in, nil }
func (r *recordLayer) WriteRecord(b []byte) (int, error)     { r.out <- b; return len(b), nil }
func (r *recordLayer) SendAlert(uint8)                       {}

func TestEncodeIntoSessionTicket(t *testing.T) {
	raddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1234}
	sIn := make(chan []byte, 10)
	sOut := make(chan []byte, 10)

	// do a first handshake and encode a "foobar" into the session ticket
	errChan := make(chan error, 1)
	stChan := make(chan []byte, 1)
	go func() {
		serverConf := testConfig.Clone()
		serverConf.AlternativeRecordLayer = &recordLayer{in: sIn, out: sOut}
		serverConf.MaxEarlyData = 1
		server := Server(&unusedConn{remoteAddr: raddr}, serverConf)
		defer server.Close()
		err := server.Handshake()
		if err != nil {
			errChan <- err
			return
		}
		st, err := server.GetSessionTicket([]byte("foobar"))
		if err != nil {
			errChan <- err
			return
		}
		stChan <- st
		errChan <- nil
	}()

	clientConf := testConfig.Clone()
	clientConf.AlternativeRecordLayer = &recordLayer{in: sOut, out: sIn}
	clientConf.ClientSessionCache = NewLRUClientSessionCache(10)
	client := Client(&unusedConn{remoteAddr: raddr}, clientConf)
	if err := client.Handshake(); err != nil {
		t.Fatalf("first handshake failed %s", err)
	}
	if err := <-errChan; err != nil {
		t.Fatalf("first handshake failed %s", err)
	}
	sOut <- <-stChan
	if err := client.HandlePostHandshakeMessage(); err != nil {
		t.Fatalf("handling the session ticket failed: %s", err)
	}
	client.Close()

	dataChan := make(chan []byte, 1)
	errChan = make(chan error, 1)
	go func() {
		serverConf := testConfig.Clone()
		serverConf.AlternativeRecordLayer = &recordLayer{in: sIn, out: sOut}
		serverConf.Accept0RTT = func(data []byte) bool {
			dataChan <- data
			return true
		}
		server := Server(&unusedConn{remoteAddr: raddr}, serverConf)
		defer server.Close()
		errChan <- server.Handshake()
	}()

	clientConf.Enable0RTT = true
	client = Client(&unusedConn{remoteAddr: raddr}, clientConf)
	if err := client.Handshake(); err != nil {
		t.Fatalf("second handshake failed %s", err)
	}
	defer client.Close()
	if err := <-errChan; err != nil {
		t.Fatalf("second handshake failed %s", err)
	}
	if len(dataChan) != 1 {
		t.Fatal("expected to receive application data")
	}
	if data := <-dataChan; !bytes.Equal(data, []byte("foobar")) {
		t.Fatalf("expected to receive a foobar, got %s", string(data))
	}
}

func TestZeroRTTRejection(t *testing.T) {
	for _, doReject := range []bool{true, false} {
		t.Run(fmt.Sprintf("doing reject: %t", doReject), func(t *testing.T) {
			raddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1234}
			sIn := make(chan []byte, 10)
			sOut := make(chan []byte, 10)

			// do a first handshake and encode a "foobar" into the session ticket
			errChan := make(chan error, 1)
			go func() {
				serverConf := testConfig.Clone()
				serverConf.AlternativeRecordLayer = &recordLayer{in: sIn, out: sOut}
				serverConf.MaxEarlyData = 1
				server := Server(&unusedConn{remoteAddr: raddr}, serverConf)
				defer server.Close()
				err := server.Handshake()
				if err != nil {
					errChan <- err
					return
				}
				st, err := server.GetSessionTicket(nil)
				if err != nil {
					errChan <- err
					return
				}
				sOut <- st
				errChan <- nil
			}()

			clientConf := testConfig.Clone()
			clientConf.AlternativeRecordLayer = &recordLayer{in: sOut, out: sIn}
			clientConf.ClientSessionCache = NewLRUClientSessionCache(10)
			client := Client(&unusedConn{remoteAddr: raddr}, clientConf)
			if err := client.Handshake(); err != nil {
				t.Fatalf("first handshake failed %s", err)
			}
			if err := <-errChan; err != nil {
				t.Fatalf("first handshake failed %s", err)
			}
			if err := client.HandlePostHandshakeMessage(); err != nil {
				t.Fatalf("handling the session ticket failed: %s", err)
			}
			client.Close()

			// now dial the second connection
			errChan = make(chan error, 1)
			connStateChan := make(chan ConnectionState, 1)
			go func() {
				serverConf := testConfig.Clone()
				serverConf.AlternativeRecordLayer = &recordLayer{in: sIn, out: sOut}
				serverConf.Accept0RTT = func(data []byte) bool { return !doReject }
				server := Server(&unusedConn{remoteAddr: raddr}, serverConf)
				defer server.Close()
				errChan <- server.Handshake()
				connStateChan <- server.ConnectionState()
			}()

			clientConf.Enable0RTT = true
			var rejected bool
			clientConf.Rejected0RTT = func() { rejected = true }
			client = Client(&unusedConn{remoteAddr: raddr}, clientConf)
			if err := client.Handshake(); err != nil {
				t.Fatalf("second handshake failed %s", err)
			}
			defer client.Close()
			if err := <-errChan; err != nil {
				t.Fatalf("second handshake failed %s", err)
			}
			if rejected != doReject {
				t.Fatal("wrong rejection")
			}
			if client.ConnectionState().Used0RTT == doReject {
				t.Fatal("wrong connection state on the client")
			}
			if (<-connStateChan).Used0RTT == doReject {
				t.Fatal("wrong connection state on the server")
			}
		})
	}
}

func TestZeroRTTRejectionOnALPNMismatch(t *testing.T) {
	raddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1234}
	sIn := make(chan []byte, 10)
	sOut := make(chan []byte, 10)

	// do a first handshake and encode a "foobar" into the session ticket
	errChan := make(chan error, 1)
	go func() {
		serverConf := testConfig.Clone()
		serverConf.NextProtos = []string{"proto1"}
		serverConf.AlternativeRecordLayer = &recordLayer{in: sIn, out: sOut}
		serverConf.MaxEarlyData = 1
		server := Server(&unusedConn{remoteAddr: raddr}, serverConf)
		defer server.Close()
		err := server.Handshake()
		if err != nil {
			errChan <- err
			return
		}
		st, err := server.GetSessionTicket(nil)
		if err != nil {
			errChan <- err
			return
		}
		sOut <- st
		errChan <- nil
	}()

	clientConf := testConfig.Clone()
	clientConf.NextProtos = []string{"proto1"}
	clientConf.AlternativeRecordLayer = &recordLayer{in: sOut, out: sIn}
	clientConf.ClientSessionCache = NewLRUClientSessionCache(10)
	client := Client(&unusedConn{remoteAddr: raddr}, clientConf)
	if err := client.Handshake(); err != nil {
		t.Fatalf("first handshake failed %s", err)
	}
	if err := <-errChan; err != nil {
		t.Fatalf("first handshake failed %s", err)
	}
	if err := client.HandlePostHandshakeMessage(); err != nil {
		t.Fatalf("handling the session ticket failed: %s", err)
	}
	client.Close()

	// now dial the second connection
	errChan = make(chan error, 1)
	connStateChan := make(chan ConnectionState, 1)
	go func() {
		serverConf := testConfig.Clone()
		serverConf.NextProtos = []string{"proto2"}
		serverConf.AlternativeRecordLayer = &recordLayer{in: sIn, out: sOut}
		serverConf.Accept0RTT = func([]byte) bool { return true }
		server := Server(&unusedConn{remoteAddr: raddr}, serverConf)
		defer server.Close()
		errChan <- server.Handshake()
		connStateChan <- server.ConnectionState()
	}()

	clientConf.Enable0RTT = true
	var rejected bool
	clientConf.Rejected0RTT = func() { rejected = true }
	clientConf.NextProtos = []string{"proto2"}
	client = Client(&unusedConn{remoteAddr: raddr}, clientConf)
	if err := client.Handshake(); err != nil {
		t.Fatalf("second handshake failed %s", err)
	}
	defer client.Close()
	if err := <-errChan; err != nil {
		t.Fatalf("second handshake failed %s", err)
	}
	if !rejected {
		t.Fatal("expected 0-RTT to be rejected")
	}
	if client.ConnectionState().Used0RTT {
		t.Fatal("expected 0-RTT to be rejected")
	}
	if (<-connStateChan).Used0RTT {
		t.Fatal("expected 0-RTT to be rejected")
	}
}
