// Package sqlserverproxy forwards TDS 7.x SQL-password sessions.
package sqlserverproxy

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"crypto/tls"
	"database/sql"
	"encoding/binary"
	"errors"
	"github.com/microsoft/go-mssqldb"
	"github.com/microsoft/go-mssqldb/msdsn"
	"linkport/internal/cluster"
	"linkport/internal/model"
	"linkport/internal/outbound"
	"linkport/internal/protocol/transport"
	"net"
	"strconv"
	"time"
)

type Handler struct{}

func New() *Handler { return &Handler{} }

type fixedDialer struct {
	dial    model.DialFunc
	address string
}

func (d fixedDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	return d.dial(ctx, "tcp", d.address)
}
func (d fixedDialer) HostName() string { host, _, _ := net.SplitHostPort(d.address); return host }
func (h *Handler) Test(ctx context.Context, p model.Proxy, dial model.DialFunc) (model.TestResult, error) {
	if err := p.CheckTargetCredential(); err != nil {
		return model.TestResult{}, err
	}
	tlsConfig, err := transport.TLSConfig(p)
	if err != nil {
		return model.TestResult{}, err
	}
	encryption := msdsn.Encryption(msdsn.EncryptionDisabled)
	if p.TargetTLS {
		encryption = msdsn.EncryptionRequired
	}
	c := mssql.NewConnectorConfig(msdsn.Config{Host: p.TargetHost, Port: uint64(p.TargetPort), User: p.TargetUser, Password: p.TargetPassword, Database: p.AuthSource, Encryption: encryption, TLSConfig: tlsConfig, Protocols: []string{"tcp"}, DisableRetry: true})
	c.Dialer = fixedDialer{dial, net.JoinHostPort(p.TargetHost, strconv.Itoa(p.TargetPort))}
	db := sql.OpenDB(c)
	defer db.Close()
	done := outbound.Begin(ctx, "protocol_session")
	err = db.PingContext(ctx)
	done(err)
	if err != nil {
		return model.TestResult{}, errors.New("SQL Server authentication or connection failed")
	}
	return model.TestResult{Version: "SQL Server TDS", TLS: p.TargetTLS}, nil
}
func (h *Handler) Serve(ctx context.Context, raw net.Conn, p model.Proxy, dial model.DialFunc, cfg *tls.Config, event model.EventFunc) error {
	defer raw.Close()
	stop := context.AfterFunc(ctx, func() { raw.Close() })
	defer stop()
	raw.SetDeadline(time.Now().Add(15 * time.Second))
	kind, hello, err := readMessage(raw)
	if err != nil {
		return err
	}
	if kind != 18 {
		return errors.New("TDS prelogin required")
	}
	enc, err := encryptionOption(hello)
	if err != nil {
		return err
	}
	var down net.Conn = raw
	if cfg != nil {
		if enc == 2 {
			return errors.New("SQL Server gateway requires TLS")
		}
		if err = writeMessage(raw, 4, prelogin(3)); err != nil {
			return err
		}
		framed := &handshakeConn{Conn: raw, framed: true}
		tlsCfg := cfg.Clone()
		tlsCfg.MaxVersion = tls.VersionTLS12
		tlsCfg.MinVersion = tls.VersionTLS12
		secure := tls.Server(framed, tlsCfg)
		if err = secure.HandshakeContext(ctx); err != nil {
			return err
		}
		framed.framed = false
		down = secure
	} else {
		if enc == 1 || enc == 3 {
			return errors.New("gateway TLS certificate required")
		}
		if err = writeMessage(raw, 4, prelogin(2)); err != nil {
			return err
		}
	}
	kind, login, err := readMessage(down)
	if err != nil {
		return err
	}
	if kind != 16 || len(login) < 94 || int(binary.LittleEndian.Uint32(login)) != len(login) || login[25]&0x80 != 0 || login[27]&1 != 0 {
		return errors.New("only SQL password LOGIN7 is supported")
	}
	user, err := loginField(login, 40, false)
	if err != nil {
		return err
	}
	password, err := loginField(login, 44, true)
	if err != nil {
		return err
	}
	account, found := transport.Account(p, user, raw.RemoteAddr())
	a, b := sha256.Sum256([]byte(password)), sha256.Sum256([]byte(account.Password))
	if !found || account.Password == "" || subtle.ConstantTimeCompare(a[:], b[:]) != 1 {
		transport.AuthenticationEvent(event, user, false)
		return errors.New("SQL Server gateway authentication failed")
	}
	p = account.Apply(p)
	dial = p.DialForTarget(dial)
	transport.AuthenticationEvent(event, user, true)
	p, _, err = cluster.Resolve(ctx, p, dial)
	if err != nil {
		return err
	}
	up, err := dial(ctx, "tcp", net.JoinHostPort(p.TargetHost, strconv.Itoa(p.TargetPort)))
	if err != nil {
		return err
	}
	defer up.Close()
	stopUp := context.AfterFunc(ctx, func() { up.Close() })
	defer stopUp()
	up.SetDeadline(time.Now().Add(15 * time.Second))
	requestEncryption := byte(2)
	if p.TargetTLS {
		requestEncryption = 1
	}
	if err = writeMessage(up, 18, prelogin(requestEncryption)); err != nil {
		return err
	}
	kind, response, err := readMessage(up)
	if err != nil {
		return err
	}
	if kind != 4 {
		return errors.New("invalid upstream prelogin")
	}
	enc, err = encryptionOption(response)
	if err != nil {
		return err
	}
	var target net.Conn = up
	if p.TargetTLS || enc == 1 || enc == 3 {
		if enc != 1 && enc != 3 {
			return errors.New("upstream TLS required")
		}
		p.TargetTLS = true
		tlsCfg, e := transport.TLSConfig(p)
		if e != nil {
			return e
		}
		tlsCfg.MaxVersion = tls.VersionTLS12
		framed := &handshakeConn{Conn: up, framed: true}
		secure := tls.Client(framed, tlsCfg)
		if err = secure.HandshakeContext(ctx); err != nil {
			return err
		}
		framed.framed = false
		target = secure
	} else if enc != 2 {
		return errors.New("unsupported upstream encryption negotiation")
	}
	login, err = upstreamLogin(login, p.TargetUser, p.TargetPassword, p.TargetHost)
	if err != nil {
		return err
	}
	if err = writeMessage(target, 16, login); err != nil {
		return err
	}
	raw.SetDeadline(time.Time{})
	up.SetDeadline(time.Time{})
	done := make(chan error, 1)
	var timing model.ExecutionQueue
	go func() { e := relayMeasured(ctx, down, target, &timing); raw.Close(); done <- e }()
	for {
		var flags byte
		kind, payload, e := readWithFlags(down, &flags)
		if e != nil {
			err = e
			break
		}
		switch kind {
		case 1:
			if len(payload) < 4 {
				err = errors.New("invalid SQL batch")
				break
			}
			offset := int(binary.LittleEndian.Uint32(payload))
			if offset < 4 || offset > len(payload) {
				err = errors.New("invalid SQL batch headers")
				break
			}
			err = model.AuditCommand(ctx, "sql_batch", decodeText(payload[offset:]), "allowed")
		case 3:
			err = model.AuditCommand(ctx, "rpc", "SQL Server RPC request (parameters omitted)", "allowed")
		case 6, 7, 14:
			err = model.AuditCommand(ctx, "tds", "SQL Server protocol operation", "allowed")
		default:
			err = errors.New("unsupported TDS operation")
		}
		if err != nil {
			break
		}
		if kind == 1 {
			timing.Start("sql_batch")
		} else if kind == 3 {
			timing.Start("rpc")
		} else {
			timing.Disable()
		}
		if err = writeWithFlags(target, kind, payload, flags); err != nil {
			break
		}
	}
	up.Close()
	raw.Close()
	<-done
	return err
}