castor/handlers/response.go

package handlers

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net"
	"net/http"
	"os"
	"path/filepath"

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

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 log interface {
	Info(...interface{})
}

func NewResponse(conn net.Conn, logger log) *Response {
	return &Response{
		logger: logger,
		conn:   conn,
	}
}

type Response struct {
	logger     log
	statusSent bool
	status     int
	meta       string
	conn       net.Conn
}

func (w *Response) 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.conn, "%d %s\r\n", code, meta)
}

func (w *Response) 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.conn.Write(b)
}

type Handler interface {
	Handle(*Response, *Request)
}

type HandlerFunc func(*Response, *Request)

// Handle calls f(w, r).
func (f HandlerFunc) Handle(w *Response, r *Request) {
	f(w, r)
}

type StaticHandler struct {
	Response string
}

func (h StaticHandler) Handle(w *Response, r *Request) {
	w.logger.Info(fmt.Sprintf("static response: %v", r.url))
	w.WriteStatus(responseCodes["SUCCESS"], "text/gemini")
	w.Write([]byte(h.Response))
}

func NewFileHandler(root string) *FileHandler {
	return &FileHandler{
		Root:             root,
		DirectoryListing: false,
		IndexFile:        "index.gmi",
	}
}

type FileHandler struct {
	Root             string
	DirectoryListing bool
	IndexFile        string
}

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)
	fmt.Fprintf(ret, "=> %s %s\r\n", filepath.Join(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 FileHandler) Handle(w *Response, r *Request) {
	// 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, h.IndexFile))
		if err == nil && sourceFileStat.Mode().IsRegular() {
			req = filepath.Join(req, h.IndexFile)
		} else if h.DirectoryListing {
			b, err := genIndex(req, filepath.Clean(r.GetURL().Path))
			if err != nil {
				w.logger.Info(fmt.Sprintf("request: %v; result: %d", r.url, responseCodes["NOTFOUND"]))
				w.WriteStatus(responseCodes["NOTFOUND"], "File not found")
				return
			}
			w.logger.Info(fmt.Sprintf("request: %v; result: %d", r.url, responseCodes["SUCCESS"]))
			w.WriteStatus(responseCodes["SUCCESS"], "text/gemini")
			w.Write(b)
			return
		}
	}

	if !sourceFileStat.Mode().IsRegular() {
		w.logger.Info(fmt.Sprintf("request: %v; result: %d", r.url, responseCodes["NOTFOUND"]))
		w.WriteStatus(responseCodes["NOTFOUND"], "File not found")
		return
	}

	source, err := os.Open(req)
	if err != nil {
		w.logger.Info(fmt.Sprintf("request: %v; result: %d", r.url, responseCodes["TEMPORARYFAILURE"]))
		w.WriteStatus(responseCodes["TEMPORARYFAILURE"], "Internal Error")
		return
	}
	defer source.Close()

	mime := mime.TypeByExtension(filepath.Ext(req))
	w.logger.Info(fmt.Sprintf("request: %v; result: %d", r.url, responseCodes["SUCCESS"]))
	w.WriteStatus(responseCodes["SUCCESS"], mime)
	io.Copy(w, source)
}

func RecoveryHandler(log logger.Logger, next Handler) Handler {
	return HandlerFunc(func(w *Response, r *Request) {
		defer func() {
			err := recover()
			if err != nil {
				log.Info(err)
				w.WriteStatus(responseCodes["TEMPORARYFAILURE"], "Internal Error")
				return
			}
		}()

		next.Handle(w, r)
	})
}