server/server.go

package main

import (
	"bufio"
	"bytes"
	"database/sql"
	"fmt"
	"os"
	"strings"

	_ "github.com/mattn/go-sqlite3"
	"github.com/valyala/fasthttp"
)

const behindProxy = true

var removeHeaders = map[string]bool{
	"date":    true,
	"expires": true,
	"server":  true,
}

var conn *sql.DB
var stmt *sql.Stmt

func main() {
	db, err := sql.Open("sqlite3", "archive.db")
	if err != nil {
		return
	}
	conn = db
	defer db.Close()

	sel_stmt, err := db.Prepare("select id, code from data where method = ? and url = ? limit 1")
	if err != nil {
		return
	}
	stmt = sel_stmt
	defer sel_stmt.Close()

	fasthttp.ListenAndServe("127.0.0.1:4001", handler)
}

func handler(ctx *fasthttp.RequestCtx) {
	// -- find in DB and read id+code
	uri := ctx.URI()

	var scheme []byte
	if behindProxy {
		scheme = ctx.Request.Header.Peek("X-Forwarded-Proto")
	} else {
		scheme = uri.Scheme()
	}

	host, port := parseHost(
		uri.Host(),
		bytes.Equal(scheme, []byte("https")),
	)

	row := stmt.QueryRow(
		string(ctx.Method()),
		fmt.Sprintf(
			"%s://%s:%s%s",
			scheme,
			host,
			port,
			ctx.RequestURI(),
		),
	)

	var id int
	var code int

	err := row.Scan(&id, &code)
	if err == sql.ErrNoRows {
		ctx.Response.SetStatusCode(404)
		return
	} else if err != nil {
		sendError(ctx, "Unable to fetch row: ", err)
		return
	}

	// -- set status code
	ctx.Response.SetStatusCode(code)

	// -- find in FS and read headers+body
	path := "storage/" + string(id)

	fh, err := os.Open(path + "/headers")
	if err != nil {
		sendError(ctx, "Unable to read headers: ", err)
		return
	}
	sc := bufio.NewScanner(fh)
	for sc.Scan() {
		header := strings.SplitN(sc.Text(), ": ", 2)
		name, value := strings.ToLower(header[0]), header[1]
		if removeHeaders[name] {
			continue
		}
		ctx.Response.Header.Add(name, value)
	}
	fh.Close()

	fb, err := os.Open(path + "/body")
	if err != nil {
		sendError(ctx, "Unable to read body: ", err)
		return
	}

	ctx.Response.SetBodyStream(fb, -1)
}

func parseHost(host []byte, https bool) ([]byte, []byte) {
	idx := bytes.LastIndex(host, []byte(":"))
	var resHost, port []byte
	if idx != -1 {
		resHost = host[:idx]
		port = host[idx+1:]
	} else {
		resHost = host
		if https {
			port = []byte("443")
		} else {
			port = []byte("80")
		}
	}
	return resHost, port
}

func sendError(ctx *fasthttp.RequestCtx, msg string, err error) {
	ctx.Response.SetStatusCode(500)
	ctx.Response.SetBodyString(msg + err.Error())
}