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
}