package checker

import (
	"bufio"
	"fmt"
	"io"
	"net"
)

// maxSTARTTLSLineBytes caps the length of a single line read from a STARTTLS
// peer. Real banners and CAPABILITY responses are well under 1 KiB; this
// bound prevents a malicious or buggy server from exhausting memory by
// withholding the line terminator.
const maxSTARTTLSLineBytes = 8 * 1024

// readLineLimited reads bytes from r up to and including the next '\n', or
// until maxSTARTTLSLineBytes have been read without one (in which case it
// returns an error). The returned string keeps the trailing '\n' so callers
// can use the same parsing logic as bufio.Reader.ReadString('\n').
func readLineLimited(r *bufio.Reader) (string, error) {
	out := make([]byte, 0, 128)
	for {
		b, err := r.ReadByte()
		if err != nil {
			if err == io.EOF && len(out) > 0 {
				return string(out), io.ErrUnexpectedEOF
			}
			return string(out), err
		}
		out = append(out, b)
		if b == '\n' {
			return string(out), nil
		}
		if len(out) >= maxSTARTTLSLineBytes {
			return string(out), fmt.Errorf("line exceeds %d bytes without terminator", maxSTARTTLSLineBytes)
		}
	}
}

// starttlsUpgrader performs the plaintext portion of a STARTTLS upgrade on
// conn, leaving conn ready for tls.Client(conn, …).Handshake(). On success
// the returned function returns nil; on failure it returns a descriptive
// error (wrap errStartTLSNotOffered when the server advertises no STARTTLS).
type starttlsUpgrader func(conn net.Conn, sni string) error

var starttlsUpgraders = map[string]starttlsUpgrader{}

func registerStartTLS(protocol string, upgrader starttlsUpgrader) {
	starttlsUpgraders[protocol] = upgrader
}

// upgraderFor returns a tlsenum-compatible upgrader callback for a given
// STARTTLS dialect, plus an ok flag. An empty dialect means direct TLS and
// returns (nil, true) — tlsenum will skip the upgrade phase. An unknown
// dialect returns (nil, false) so the caller can record the skip reason.
func upgraderFor(dialect, sni string) (func(net.Conn) error, bool) {
	if dialect == "" {
		return nil, true
	}
	up, ok := starttlsUpgraders[dialect]
	if !ok {
		return nil, false
	}
	return func(c net.Conn) error { return up(c, sni) }, true
}