// SPDX-License-Identifier: AGPL-3.0-or-later

package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	_ "embed"
	"encoding/binary"
	"encoding/json"
	"flag"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"strings"
	"time"
)

const readTimeout = 10  // seconds
const writeTimeout = 10 // seconds

type tlsConnectionInfo struct {
	TlsVersion         uint16 `json:"tls_version"`
	CipherSuite        uint16 `json:"cipher_suite"`
	SessionResumed     bool   `json:"session_resumed"`
	NegotiatedProtocol string `json:"alpn_negotiated_protocol"` // ALPN
}

// Connection wrapper that enables exposing the Client Hello to the
// request handler.
// See https://github.com/FiloSottile/mostly-harmless/tree/main/talks/asyncnet
type prefixConn struct {
	net.Conn
	io.Reader
}

func (c prefixConn) Read(p []byte) (int, error) {
	return c.Reader.Read(p)
}

// Output to stderr and exit with error code 1.
// Like log.Fatal, but without the date&time prefix.
// Used before starting the server loop.
func fatalError(err ...any) {
	logger := log.New(os.Stderr, "", 0)
	logger.Fatal(err...)
}

//go:embed index.html
var html string

//go:embed index.gmi
var gemtext string

// Copy the Client Hello message before starting the TLS handshake.
func peek(conn net.Conn, tlsConfig *tls.Config) {
	defer conn.Close()
	var buf bytes.Buffer
	// Copy TLS record header.
	_, err := io.CopyN(&buf, conn, 5)
	if err != nil {
		log.Println(err)
		return
	}
	// Check if this is a TLS handshake record.
	if buf.Bytes()[0] != 0x16 {
		return
	}
	// Extract handshake message length.
	handshakeMessageLength := binary.BigEndian.Uint16(buf.Bytes()[3:5])
	// Copy handshake message (should be a Client Hello).
	_, err = io.CopyN(&buf, conn, int64(handshakeMessageLength))
	if err != nil {
		log.Println(err)
		return
	}
	rawClientHello := buf.Bytes()

	// "Put back" the Client Hello bytes we just read, so that they can be
	// used in the TLS handshake.  Concatenate the read bytes with the
	// unread bytes using a MultiReader, inside a connection wrapper.
	pConn := prefixConn{
		Conn:   conn,
		Reader: io.MultiReader(&buf, conn),
	}

	tlsConnection := tls.Server(pConn, tlsConfig)
	err = tlsConnection.Handshake()
	if err != nil {
		log.Println(err)
		return
	}
	requestHandler(tlsConnection, rawClientHello)
}

func requestHandler(conn *tls.Conn, rawClientHello []byte) {
	defer conn.Close()

	scanner := bufio.NewScanner(conn)
	scanner.Scan()
	line := scanner.Text()
	err := scanner.Err()
	if err != nil {
		log.Println(err)
		return
	}

	var protocol string
	var path string // requested page

	// Detect protocol by checking the first few characters.
	if strings.HasPrefix(line, "gemini://") {
		protocol = "Gemini"
		u, err := url.Parse(line)
		if err != nil {
			log.Println(err)
			return
		}
		path = u.Path
		if path == "" {
			path = "/"
		}
		if path == "/" {
			_, err = conn.Write([]byte("20 text/gemini\r\n"))
		} else if path == "/json/v1" {
			_, err = conn.Write([]byte("20 application/json\r\n"))
		}
	} else if strings.HasPrefix(line, "GET ") ||
		strings.HasPrefix(line, "POST ") ||
		strings.HasPrefix(line, "HEAD ") {
		protocol = "HTTP"
		path = strings.Split(line, " ")[1]
		if path == "/" {
			_, err = conn.Write([]byte("HTTP/1.1 200 OK\r\n"))
			_, err = conn.Write([]byte("Content-Type: text/html; charset=utf-8\r\n"))
			_, err = conn.Write([]byte("Cache-Control: no-store\r\n"))
			_, err = conn.Write([]byte("Vary: *\r\n"))
		} else if path == "/json/v1" {
			_, err = conn.Write([]byte("HTTP/1.1 200 OK\r\n"))
			_, err = conn.Write([]byte("Content-Type: application/json\r\n"))
			_, err = conn.Write([]byte("Cache-Control: no-store\r\n"))
			_, err = conn.Write([]byte("Vary: *\r\n"))
		} else {
			_, err = conn.Write([]byte("HTTP/1.1 404 Not Found\r\n"))
		}
		_, err = conn.Write([]byte("\r\n")) // end headers
		if strings.HasPrefix(line, "HEAD ") {
			return // headers only
		}
	} else {
		log.Println("Unknown protocol")
		return
	}

	var clientHello clientHelloMsg
	if !clientHello.unmarshal(rawClientHello) {
		log.Println("Failed to parse Client Hello")
		return
	}

	connectionState := conn.ConnectionState() // TLS
	tlsConnInfo := tlsConnectionInfo{
		connectionState.Version,
		connectionState.CipherSuite,
		connectionState.DidResume,
		connectionState.NegotiatedProtocol,
	}

	if path == "/" {
		if protocol == "HTTP" {
			_, err = conn.Write([]byte(html))
		} else if protocol == "Gemini" {
			_, err = conn.Write([]byte(gemtext))
		} else {
			log.Println("Unknown protocol")
			return
		}
	} else if path == "/json/v1" {
		output := struct {
			ClientHello       clientHelloMsg    `json:"client_hello"`
			TlsConnectionInfo tlsConnectionInfo `json:"connection_info"`
		}{
			clientHello,
			tlsConnInfo,
		}
		outputJSON, err := json.MarshalIndent(output, "", "  ")
		if err != nil {
			log.Println(err)
			return
		}
		_, err = conn.Write(outputJSON)
	}
	if err != nil {
		log.Println(err)
		return
	}
}

func main() {

	var certFile, keyFile string
	var userToSwitchTo string
	var hostAndPort string

	// Parse arguments
	flag.StringVar(&certFile, "c", "", "path to certificate file")
	flag.StringVar(&keyFile, "k", "", "path to private key file")
	flag.StringVar(&userToSwitchTo, "u", "", "user to switch to, if running as root")
	flag.Parse()
	hostAndPort = flag.Arg(0)
	if certFile == "" || keyFile == "" || hostAndPort == "" {
		fatalError("usage: client-hello-mirror -c cert.pem -k key.pem [-u user] host:port")
	}

	// Load cert
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		fatalError(err)
	}
	// TLS config
	tlsConfig := tls.Config{
		Certificates: []tls.Certificate{cert},
		MinVersion:   tls.VersionTLS10,
		NextProtos:   []string{"http/1.1"},
	}

	// Listen for connections
	ln, err := net.Listen("tcp", hostAndPort)
	if err != nil {
		fatalError(err)
	}
	defer ln.Close()

	dropPrivileges(userToSwitchTo)

	log.Println("Server started")

	for {
		// Wait for a connection.
		conn, err := ln.Accept()
		if err != nil {
			log.Println("Error accepting: ", err.Error())
			continue
		}
		// Set connection timeouts.
		err = conn.SetReadDeadline(time.Now().Add(readTimeout * time.Second))
		if err != nil {
			log.Println("SetReadDeadline error: ", err.Error())
			continue
		}
		err = conn.SetWriteDeadline(time.Now().Add(writeTimeout * time.Second))
		if err != nil {
			log.Println("SetWriteDeadline error: ", err.Error())
			continue
		}
		// Process the request.
		go peek(conn, &tlsConfig)
	}
}