package tcp

import (
	"github.com/apernet/OpenGFW/analyzer"
	"github.com/apernet/OpenGFW/analyzer/internal"
	"github.com/apernet/OpenGFW/analyzer/utils"
)

var _ analyzer.TCPAnalyzer = (*TLSAnalyzer)(nil)

type TLSAnalyzer struct{}

func (a *TLSAnalyzer) Name() string {
	return "tls"
}

func (a *TLSAnalyzer) Limit() int {
	return 8192
}

func (a *TLSAnalyzer) NewTCP(info analyzer.TCPInfo, logger analyzer.Logger) analyzer.TCPStream {
	return newTLSStream(logger)
}

const (
	// tlsRecordHeaderSize is content type (1 byte) + legacy protocol version
	// (2 bytes) + length (2 bytes).
	tlsRecordHeaderSize = 5
	// tlsMaxPlaintextLen is the largest body a plaintext record may carry.
	tlsMaxPlaintextLen = 1 << 14
	// tlsHandshakeHeaderSize is message type (1 byte) + length (3 bytes).
	tlsHandshakeHeaderSize = 4
)

type tlsStream struct {
	logger analyzer.Logger

	req  tlsHello
	resp tlsHello
}

// tlsHello reads one hello message off one direction of the stream.
//
// A handshake message may be split across any number of records (RFC 8446
// 5.1), and splitting the ClientHello into tiny records is a common way to
// hide the SNI from middleboxes that expect it in the first record. So the
// record layer is peeled off first and the message is rebuilt from the record
// bodies before it is parsed.
type tlsHello struct {
	rec utils.ByteBuffer // record layer bytes not consumed yet
	msg utils.ByteBuffer // handshake layer bytes collected so far
	lsm *utils.LinearStateMachine

	msgType byte
	minLen  int                                      // smallest valid message body
	parse   func(*utils.ByteBuffer) analyzer.PropMap // parses the message body

	recLeft  int // body bytes of the current record not consumed yet
	msgLen   int // message body length, -1 until the handshake header is read
	records  int // records that have carried the message so far
	reported int // value of records last published

	m       analyzer.PropMap
	updated bool
	done    bool
}

func newTLSStream(logger analyzer.Logger) *tlsStream {
	s := &tlsStream{logger: logger}
	// minimal ClientHello body: protocol version (2 bytes) + random (32 bytes) +
	//   + session ID (1 byte) + cipher suites (4 bytes) +
	//   + compression methods (2 bytes) + no extensions
	s.req.init(internal.TypeClientHello, 41, internal.ParseTLSClientHelloMsgData)
	// minimal ServerHello body: server version (2 bytes) + random (32 bytes) +
	//   + session ID (>=1 byte) + cipher suite (2 bytes) +
	//   + compression method (1 byte) + no extensions
	s.resp.init(internal.TypeServerHello, 38, internal.ParseTLSServerHelloMsgData)
	return s
}

func (h *tlsHello) init(msgType byte, minLen int, parse func(*utils.ByteBuffer) analyzer.PropMap) {
	h.msgType, h.minLen, h.parse = msgType, minLen, parse
	h.msgLen = -1
	h.lsm = utils.NewLinearStateMachine(h.readMessage, h.parseMessage)
}

func (s *tlsStream) Feed(rev, start, end bool, skip int, data []byte) (u *analyzer.PropUpdate, done bool) {
	if skip != 0 {
		return nil, true
	}
	if len(data) == 0 {
		return nil, false
	}
	h, key := &s.req, "req"
	if rev {
		h, key = &s.resp, "resp"
	}
	if h.done {
		return nil, false
	}
	h.rec.Append(data)
	h.updated = false
	var cancelled bool
	cancelled, h.done = h.lsm.Run()
	// A fragmented hello is published as soon as it shows up, not only once
	// the message is complete: a peer can make it too large to ever complete
	// within Limit(), and a rule should not need the SNI to act on it.
	if m := h.update(key, h.updated || (h.records > 1 && h.reported <= 1)); m != nil {
		u = &analyzer.PropUpdate{Type: analyzer.PropUpdateMerge, M: m}
	}
	return u, cancelled || (s.req.done && s.resp.done)
}

// update returns the properties to publish for this direction, or nil if
// there is nothing to publish. The parsed message goes under key, the record
// count under key + "_records". The record count is kept out of the message
// map so that a non-nil message map still means a fully parsed hello.
func (h *tlsHello) update(key string, publish bool) analyzer.PropMap {
	if !publish || h.records == 0 {
		return nil
	}
	m := analyzer.PropMap{key + "_records": h.records}
	if h.updated {
		m[key] = h.m
		h.updated = false
	}
	h.reported = h.records
	return m
}

// readMessage consumes records until a whole hello message is collected in
// h.msg.
func (h *tlsHello) readMessage() utils.LSMAction {
	for {
		if h.msgLen < 0 && h.msg.Len() >= tlsHandshakeHeaderSize {
			header := h.msg.Buf
			if header[0] != h.msgType {
				return utils.LSMActionCancel
			}
			h.msgLen = int(header[1])<<16 | int(header[2])<<8 | int(header[3])
			if h.msgLen < h.minLen {
				return utils.LSMActionCancel
			}
		}
		if h.msgLen >= 0 && h.msg.Len() >= tlsHandshakeHeaderSize+h.msgLen {
			return utils.LSMActionNext
		}
		if h.recLeft == 0 {
			header, ok := h.rec.Get(tlsRecordHeaderSize, true)
			if !ok {
				return utils.LSMActionPause
			}
			if header[0] != internal.RecordTypeHandshake {
				// Nothing else may come between the records of a handshake
				// message.
				return utils.LSMActionCancel
			}
			h.recLeft = int(header[3])<<8 | int(header[4])
			if h.recLeft > tlsMaxPlaintextLen {
				return utils.LSMActionCancel
			}
			// Empty records are forbidden for handshake messages, but still
			// counted rather than rejected: a peer that sends them anyway is
			// fragmenting, and rejecting them would stop the analysis.
			h.records++
			continue
		}
		n := min(h.recLeft, h.rec.Len())
		if n == 0 {
			return utils.LSMActionPause
		}
		body, _ := h.rec.Get(n, true)
		h.recLeft -= n
		if h.msg.Len() == 0 {
			// Usually the whole message sits in this one chunk, so reference
			// it instead of copying. The capacity is capped so that a later
			// Append reallocates rather than writing into h.rec.
			h.msg.Buf = body[:n:n]
		} else {
			h.msg.Append(body)
		}
	}
}

// parseMessage converts the collected hello message into h.m.
//
// Anything after the message, such as the rest of the server's flight, is
// ignored.
func (h *tlsHello) parseMessage() utils.LSMAction {
	body := &utils.ByteBuffer{Buf: h.msg.Buf[tlsHandshakeHeaderSize : tlsHandshakeHeaderSize+h.msgLen]}
	m := h.parse(body)
	h.rec.Reset()
	h.msg.Reset()
	if m == nil {
		return utils.LSMActionCancel
	}
	h.m = m
	h.updated = true
	return utils.LSMActionNext
}

func (s *tlsStream) Close(limited bool) *analyzer.PropUpdate {
	// Report how far a fragmented hello got if it never completed, e.g.
	// because it was stretched past Limit().
	m := s.req.update("req", s.req.records > 1 && s.req.records > s.req.reported)
	if rm := s.resp.update("resp", s.resp.records > 1 && s.resp.records > s.resp.reported); rm != nil {
		if m == nil {
			m = rm
		} else {
			for k, v := range rm {
				m[k] = v
			}
		}
	}
	s.req.close()
	s.resp.close()
	if m == nil {
		return nil
	}
	return &analyzer.PropUpdate{Type: analyzer.PropUpdateMerge, M: m}
}

func (h *tlsHello) close() {
	h.rec.Reset()
	h.msg.Reset()
	h.m = nil
}