package sqlserverproxy

import (
	"context"
	"crypto/tls"
	"database/sql"
	"encoding/hex"
	"encoding/pem"
	"github.com/microsoft/go-mssqldb"
	"github.com/microsoft/go-mssqldb/msdsn"
	"linkport/internal/model"
	"net"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"
)

func TestDriverThroughGateway(t *testing.T) {
	for _, encrypted := range []bool{false, true} {
		t.Run(strconv.FormatBool(encrypted), func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			certServer := httptest.NewTLSServer(nil)
			cert := certServer.TLS.Certificates[0]
			certServer.Close()
			up, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatal(err)
			}
			defer up.Close()
			captured := make(chan []string, 1)
			upDone := make(chan struct{})
			go func() {
				defer close(upDone)
				raw, e := up.Accept()
				if e != nil {
					return
				}
				defer raw.Close()
				raw.SetDeadline(time.Now().Add(8 * time.Second))
				if _, _, e = readMessage(raw); e != nil {
					t.Error(e)
					return
				}
				var conn net.Conn = raw
				enc := byte(2)
				if encrypted {
					enc = 1
				}
				if e = writeMessage(raw, 4, prelogin(enc)); e != nil {
					t.Error(e)
					return
				}
				if encrypted {
					framed := &handshakeConn{Conn: raw, framed: true}
					secure := tls.Server(framed, &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12, MaxVersion: tls.VersionTLS12})
					if e = secure.HandshakeContext(ctx); e != nil {
						t.Error(e)
						return
					}
					framed.framed = false
					conn = secure
				}
				kind, login, e := readMessage(conn)
				if e != nil || kind != 16 {
					t.Errorf("login: %d %v", kind, e)
					return
				}
				user, e := loginField(login, 40, false)
				if e != nil {
					t.Error(e)
					return
				}
				password, e := loginField(login, 44, true)
				if e != nil {
					t.Error(e)
					return
				}
				captured <- []string{user, password}
				ack, _ := hex.DecodeString("ad0a000174000004000f000000fd000000000000000000000000")
				if e = writeMessage(conn, 4, ack); e != nil {
					t.Error(e)
					return
				}
				for {
					kind, _, e = readMessage(conn)
					if e != nil {
						return
					}
					if kind != 1 && kind != 3 {
						t.Errorf("unexpected request %d", kind)
						return
					}
					if e = writeMessage(conn, 4, []byte{253, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}); e != nil {
						return
					}
				}
			}()
			host, port, _ := net.SplitHostPort(up.Addr().String())
			number, _ := strconv.Atoi(port)
			p := model.Proxy{Type: "sqlserver", TargetHost: host, TargetPort: number, TargetTLS: encrypted, TargetCA: string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: cert.Certificate[0]})), GatewayAccounts: []model.GatewayAccount{{ID: "client", Username: "client", Password: "client-secret", TargetAccountID: "mapped"}}, TargetAccounts: []model.TargetAccount{{ID: "mapped", Username: "upstream", Password: "upstream-secret"}}}
			gateway, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatal(err)
			}
			defer gateway.Close()
			done := make(chan error, 1)
			var gatewayTLS *tls.Config
			if encrypted {
				gatewayTLS = &tls.Config{Certificates: []tls.Certificate{cert}}
			}
			go func() {
				raw, e := gateway.Accept()
				if e != nil {
					done <- e
					return
				}
				done <- New().Serve(ctx, raw, p, (&net.Dialer{}).DialContext, gatewayTLS, nil)
			}()
			_, gatewayPort, _ := net.SplitHostPort(gateway.Addr().String())
			n, _ := strconv.Atoi(gatewayPort)
			encryption := msdsn.Encryption(msdsn.EncryptionDisabled)
			if encrypted {
				encryption = msdsn.EncryptionRequired
			}
			connector := mssql.NewConnectorConfig(msdsn.Config{Host: "127.0.0.1", Port: uint64(n), User: "client", Password: "client-secret", Encryption: encryption, TLSConfig: &tls.Config{InsecureSkipVerify: true}, Protocols: []string{"tcp"}})
			db := sql.OpenDB(connector)
			defer db.Close()
			if err = db.PingContext(ctx); err != nil {
				t.Fatal(err)
			}
			select {
			case credentials := <-captured:
				if credentials[0] != "upstream" || credentials[1] != "upstream-secret" {
					t.Fatal("wrong mapped credentials")
				}
			case <-ctx.Done():
				t.Fatal(ctx.Err())
			}
			db.Close()
			cancel()
			select {
			case <-done:
			case <-time.After(2 * time.Second):
				t.Fatal("gateway did not close")
			}
			select {
			case <-upDone:
			case <-time.After(2 * time.Second):
				t.Fatal("upstream did not close")
			}
		})
	}
}