castor/main.go

package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"

	"git.yetaga.in/alazyreader/castor/certificates"
)

var responseCodes = map[string]int{
	"INPUT":                     10,
	"SENSITIVEINPUT":            11,
	"SUCCESS":                   20,
	"REDIRECTTEMPORARY":         30,
	"REDIRECTPERMANENT":         31,
	"TEMPORARYFAILURE":          40,
	"SERVERUNAVAILABLE":         41,
	"CGIERROR":                  42,
	"PROXYERROR":                43,
	"SLOWDOWN":                  44,
	"PERMANENTFAILURE":          50,
	"NOTFOUND":                  51,
	"GONE":                      52,
	"PROXYREQUESTREFUSED":       53,
	"BADREQUEST":                59,
	"CLIENTCERTIFICATEREQUIRED": 60,
	"CERTIFICATENOTAUTHORISED":  61,
	"CERTIFICATENOTVALID":       62,
}

type geminiRequest struct {
	url *url.URL
}

func (r *geminiRequest) GetURL() *url.URL {
	return r.url
}

type geminiResponse struct {
	statusSent bool
	status     int
	meta       string
	connection net.Conn
}

func (w *geminiResponse) WriteStatus(code int, meta string) (int, error) {
	if w.statusSent {
		return 0, fmt.Errorf("Cannot set status after start of response")
	}
	w.status = code
	w.meta = meta
	w.statusSent = true
	return fmt.Fprintf(w.connection, "%d %s\r\n", code, meta)
}

func (w *geminiResponse) Write(b []byte) (int, error) {
	if !w.statusSent {
		// this can't guess text/gemini, of course.
		guessedType := http.DetectContentType(b)
		w.WriteStatus(responseCodes["SUCCESS"], guessedType)
	}
	return w.connection.Write(b)
}

type geminiHandler interface {
	Handle(*geminiResponse, *geminiRequest)
}

type geminiHandlerFunc func(*geminiResponse, *geminiRequest)

// Handle calls f(w, r).
func (f geminiHandlerFunc) Handle(w *geminiResponse, r *geminiRequest) {
	f(w, r)
}

// handlers
type staticGeminiHandler struct {
	StaticString string
}

func (h staticGeminiHandler) Handle(w *geminiResponse, r *geminiRequest) {
	w.Write([]byte(h.StaticString))
}

type fsGeminiHandler struct {
	root             string
	DirectoryListing bool
}

func genIndex(folder, rel string) ([]byte, error) {
	files, err := ioutil.ReadDir(folder)
	if err != nil {
		return []byte{}, err
	}
	ret := bytes.NewBuffer([]byte{})

	fmt.Fprintf(ret, "# %s\r\n\r\n", rel)
	for _, file := range files {
		fmt.Fprintf(ret, "=> %s %s\r\n", filepath.Join(rel, file.Name()), file.Name())
	}
	return ret.Bytes(), nil
}

func (h fsGeminiHandler) Handle(w *geminiResponse, r *geminiRequest) {
	// Clean, then join; can't escape the defined root
	req := filepath.Join(h.root, filepath.Clean(r.GetURL().Path))

	sourceFileStat, err := os.Stat(req)
	if err != nil {
		w.WriteStatus(responseCodes["NOTFOUND"], "File not found")
		return
	}

	if sourceFileStat.IsDir() {
		sourceFileStat, err = os.Stat(filepath.Join(req, "index.gmi"))
		if err == nil && sourceFileStat.Mode().IsRegular() {
			// if it's a directory, transparently insert the index.gmi check
			req = filepath.Join(req, "index.gmi")
		} else if h.DirectoryListing {
			b, err := genIndex(req, filepath.Clean(r.GetURL().Path))
			if err != nil {
				w.WriteStatus(responseCodes["NOTFOUND"], "File not found")
				return
			}
			w.WriteStatus(responseCodes["SUCCESS"], "text/gemini")
			w.Write(b)
			return
		}
	}

	if !sourceFileStat.Mode().IsRegular() {
		w.WriteStatus(responseCodes["NOTFOUND"], "File not found")
		return
	}

	source, err := os.Open(req)
	if err != nil {
		w.WriteStatus(responseCodes["TEMPORARYFAILURE"], "Internal Error")
		return
	}
	defer source.Close()

	mime := mime.TypeByExtension(filepath.Ext(req))
	w.WriteStatus(responseCodes["SUCCESS"], mime)
	io.Copy(w, source)
}

func recoveryHandler(next geminiHandler) geminiHandler {
	return geminiHandlerFunc(func(w *geminiResponse, r *geminiRequest) {
		defer func() {
			err := recover()
			if err != nil {
				log.Println(err)
				w.WriteStatus(responseCodes["TEMPORARYFAILURE"], "Internal Error")
				return
			}
		}()

		next.Handle(w, r)
	})
}

func handleConnection(log Logger, conn net.Conn, h geminiHandler) {
	defer conn.Close()
	scanner := bufio.NewScanner(conn)
	if ok := scanner.Scan(); !ok {
		log.Info(scanner.Err())
	}
	u, err := url.Parse(scanner.Text())
	if err != nil {
		log.Info(err)
	}
	w := &geminiResponse{
		connection: conn,
	}
	r := &geminiRequest{
		url: u,
	}
	recoveryHandler(h).Handle(w, r)
}

func main() {
	log := NewLogger(true)

	mime.AddExtensionType(".gemini", "text/gemini")
	mime.AddExtensionType(".gmi", "text/gemini")

	err := certificates.TestCertificateExists("./cert.pem", "./key.pem")
	var cer tls.Certificate
	if err != nil {
		log.Info("Generating new certificate...")
		key, cert, err := certificates.GenerateKeyPair("localhost")
		if err != nil {
			log.Info("error generating certificates", err)
			return
		}
		err = certificates.WriteCertsToFile("./cert.pem", "./key.pem", cert, key)
		if err != nil {
			log.Info("error saving certificates", err)
			return
		}
	}

	cer, err = tls.LoadX509KeyPair("./cert.pem", "./key.pem")
	if err != nil {
		log.Info("error loading certificates", err)
		return
	}

	tlsc := tls.Config{
		Certificates: []tls.Certificate{cer},
		MinVersion:   tls.VersionTLS12,
	}

	listener, err := tls.Listen("tcp", "localhost:1965", &tlsc)
	if err != nil {
		log.Info(err)
		return
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Debug(err.Error())
			continue
		}

		go handleConnection(log, conn, fsGeminiHandler{
			root:             "./root/",
			DirectoryListing: true,
		})
	}
}