twins-upstream/server.go

package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/h2non/filetype"
	"github.com/makeworld-the-better-one/go-gemini"
)

const readTimeout = 30 * time.Second

func writeHeader(c net.Conn, code int, meta string) {
	fmt.Fprintf(c, "%d %s\r\n", code, meta)

	if verbose {
		log.Printf("< %d %s\n", code, meta)
	}
}

func writeStatus(c net.Conn, code int) {
	var meta string
	switch code {
	case gemini.StatusTemporaryFailure:
		meta = "Temporary failure"
	case gemini.StatusProxyError:
		meta = "Proxy error"
	case gemini.StatusBadRequest:
		meta = "Bad request"
	case gemini.StatusNotFound:
		meta = "Not found"
	case gemini.StatusProxyRequestRefused:
		meta = "Proxy request refused"
	}
	writeHeader(c, code, meta)
}

func serveFile(c net.Conn, filePath string) {
	fi, err := os.Stat(filePath)
	if os.IsNotExist(err) {
		writeStatus(c, gemini.StatusNotFound)
		return
	} else if err != nil {
		writeStatus(c, gemini.StatusTemporaryFailure)
		return
	}

	if mode := fi.Mode(); mode.IsDir() {
		_, err := os.Stat(path.Join(filePath, "index.gemini"))
		if err == nil {
			filePath = path.Join(filePath, "index.gemini")
		} else {
			filePath = path.Join(filePath, "index.gmi")
		}
	}

	fi, err = os.Stat(filePath)
	if os.IsNotExist(err) {
		writeStatus(c, gemini.StatusNotFound)
		return
	} else if err != nil {
		writeStatus(c, gemini.StatusTemporaryFailure)
		return
	}

	// Open file
	file, _ := os.Open(filePath)
	defer file.Close()

	// Read file header
	buf := make([]byte, 261)
	n, _ := file.Read(buf)

	// Write header
	var mimeType string
	if strings.HasSuffix(filePath, ".html") && strings.HasSuffix(filePath, ".htm") {
		mimeType = "text/html; charset=utf-8"
	} else if strings.HasSuffix(filePath, ".txt") && strings.HasSuffix(filePath, ".text") {
		mimeType = "text/plain; charset=utf-8"
	} else if !strings.HasSuffix(filePath, ".gmi") && !strings.HasSuffix(filePath, ".gemini") {
		kind, _ := filetype.Match(buf[:n])
		if kind != filetype.Unknown {
			mimeType = kind.MIME.Value
		}
	}
	if mimeType == "" {
		mimeType = "text/gemini; charset=utf-8"
	}
	writeHeader(c, gemini.StatusSuccess, mimeType)

	// Write body
	c.Write(buf[:n])
	io.Copy(c, file)
}

func serveProxy(c net.Conn, requestData, proxyURL string) {
	original := proxyURL

	tlsConfig := &tls.Config{}
	if strings.HasPrefix(proxyURL, "gemini://") {
		proxyURL = proxyURL[9:]
	} else if strings.HasPrefix(proxyURL, "gemini-insecure://") {
		proxyURL = proxyURL[18:]
		tlsConfig.InsecureSkipVerify = true
	}
	proxy, err := tls.Dial("tcp", proxyURL, tlsConfig)
	if err != nil {
		writeStatus(c, gemini.StatusProxyError)
		return
	}
	defer proxy.Close()

	// Forward request
	proxy.Write([]byte(requestData))
	proxy.Write([]byte("\r\n"))

	// Forward response
	io.Copy(c, proxy)

	if verbose {
		log.Printf("< %s\n", original)
	}
}

func serveCommand(c net.Conn, command []string) {
	var args []string
	if len(command) > 0 {
		args = command[1:]
	}
	cmd := exec.Command(command[0], args...)

	var buf bytes.Buffer
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	err := cmd.Run()
	if err != nil {
		writeStatus(c, gemini.StatusProxyError)
		return
	}

	writeHeader(c, gemini.StatusSuccess, "text/gemini; charset=utf-8")
	c.Write(buf.Bytes())

	if verbose {
		log.Printf("< %s\n", command)
	}
}

func scanCRLF(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexByte(data, '\r'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, data[0:i], nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), data, nil
	}
	// Request more data.
	return 0, nil, nil
}

func handleConn(c net.Conn) {
	if verbose {
		t := time.Now()
		defer func() {
			d := time.Since(t)
			if d > time.Second {
				d = d.Round(time.Second)
			} else {
				d = d.Round(time.Millisecond)
			}
			log.Printf("took %s", d)
		}()
	}

	defer c.Close()
	c.SetReadDeadline(time.Now().Add(readTimeout))

	var requestData string
	scanner := bufio.NewScanner(c)
	scanner.Split(scanCRLF)
	if scanner.Scan() {
		requestData = scanner.Text()
	}
	if err := scanner.Err(); err != nil {
		writeStatus(c, gemini.StatusBadRequest)
		return
	}

	if verbose {
		log.Printf("> %s\n", requestData)
	}

	if len(requestData) > 1024 || !utf8.ValidString(requestData) {
		writeStatus(c, gemini.StatusBadRequest)
		return
	}

	request, err := url.Parse(requestData)
	if err != nil || request.Hostname() == "" || strings.ContainsRune(request.Hostname(), ' ') {
		writeStatus(c, gemini.StatusBadRequest)
		return
	}

	var requestPort int
	if request.Port() != "" {
		requestPort, err = strconv.Atoi(request.Port())
		if err != nil {
			requestPort = 0
		}
	}

	if request.Scheme == "" {
		request.Scheme = "gemini"
	}
	if request.Scheme != "gemini" || request.Hostname() != config.Hostname || (requestPort > 0 && requestPort != config.Port) {
		writeStatus(c, gemini.StatusProxyRequestRefused)
	}

	if request.Path == "" {
		// Redirect to /
		writeHeader(c, gemini.StatusRedirectPermanent, requestData+"/")
		return
	}

	pathBytes := []byte(request.Path)
	strippedPath := request.Path
	if strippedPath[0] == '/' {
		strippedPath = strippedPath[1:]
	}

	for _, serve := range config.Serve {
		if serve.Proxy != "" {
			if serve.r != nil && serve.r.Match(pathBytes) {
				serveProxy(c, requestData, serve.Proxy)
				return
			} else if serve.r == nil && strings.HasPrefix(request.Path, serve.Path) {
				serveProxy(c, requestData, serve.Proxy)
				return
			}
		} else if serve.cmd != nil {
			if serve.r != nil && serve.r.Match(pathBytes) {
				serveCommand(c, serve.cmd)
				return
			} else if serve.r == nil && strings.HasPrefix(request.Path, serve.Path) {
				serveCommand(c, serve.cmd)
				return
			}
		}

		if serve.r != nil && serve.r.Match(pathBytes) {
			serveFile(c, path.Join(serve.Root, strippedPath))
			return
		} else if serve.r == nil && strings.HasPrefix(request.Path, serve.Path) {
			serveFile(c, path.Join(serve.Root, strippedPath[len(serve.Path)-1:]))
			return
		}
	}

	writeStatus(c, gemini.StatusNotFound)
}

func handleListener(l net.Listener) {
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Fatal(err)
		}

		go handleConn(conn)
	}
}

func listen(hostname string, port int, certFile, keyFile string) {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Fatalf("failed to load certificate: %s", err)
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	listener, err := tls.Listen("tcp", fmt.Sprintf("%s:%d", hostname, port), tlsConfig)
	if err != nil {
		log.Fatalf("failed to listen on %s:%d: %s", hostname, port, err)
	}

	handleListener(listener)
}