feat: add config via args, log ListenAndServe error if any
This commit is contained in:
parent
94897e7711
commit
78305a063a
1 changed files with 25 additions and 8 deletions
33
server.go
33
server.go
|
@ -4,7 +4,9 @@ import (
|
|||
"bufio"
|
||||
"bytes"
|
||||
"database/sql"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
@ -13,7 +15,15 @@ import (
|
|||
"github.com/valyala/fasthttp"
|
||||
)
|
||||
|
||||
const behindProxy = true
|
||||
const pathsep = string(os.PathSeparator)
|
||||
|
||||
var (
|
||||
bind = flag.String("bind", "127.0.0.1:4001", "Where to listen for connections")
|
||||
proxy = flag.Bool("proxy", true, "Whether server is behind a reverse proxy")
|
||||
verbose = flag.Bool("verbose", false, "Whether to log all URLs")
|
||||
dbpath = flag.String("dbpath", "archive.db", "Path to SQLite3 DB generated by mitmproxy addon")
|
||||
storage = flag.String("storage", "storage", "Path to archived responses storage")
|
||||
)
|
||||
|
||||
var removeHeaders = map[string]bool{
|
||||
"date": true,
|
||||
|
@ -25,7 +35,9 @@ var conn *sql.DB
|
|||
var stmt *sql.Stmt
|
||||
|
||||
func main() {
|
||||
db, err := sql.Open("sqlite3", "archive.db")
|
||||
flag.Parse()
|
||||
|
||||
db, err := sql.Open("sqlite3", *dbpath)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
@ -39,7 +51,10 @@ func main() {
|
|||
stmt = sel_stmt
|
||||
defer sel_stmt.Close()
|
||||
|
||||
fasthttp.ListenAndServe("127.0.0.1:4001", handler)
|
||||
err1 := fasthttp.ListenAndServe(*bind, handler)
|
||||
if err1 != nil {
|
||||
log.Fatalln(err1)
|
||||
}
|
||||
}
|
||||
|
||||
func handler(ctx *fasthttp.RequestCtx) {
|
||||
|
@ -47,7 +62,7 @@ func handler(ctx *fasthttp.RequestCtx) {
|
|||
uri := ctx.URI()
|
||||
|
||||
var scheme []byte
|
||||
if behindProxy {
|
||||
if *proxy {
|
||||
scheme = ctx.Request.Header.Peek("X-Forwarded-Proto")
|
||||
} else {
|
||||
scheme = uri.Scheme()
|
||||
|
@ -70,7 +85,9 @@ func handler(ctx *fasthttp.RequestCtx) {
|
|||
urlToSearch,
|
||||
)
|
||||
|
||||
// log.Println(urlToSearch)
|
||||
if *verbose {
|
||||
log.Println(urlToSearch)
|
||||
}
|
||||
|
||||
var id int
|
||||
var code int
|
||||
|
@ -88,9 +105,9 @@ func handler(ctx *fasthttp.RequestCtx) {
|
|||
ctx.Response.SetStatusCode(code)
|
||||
|
||||
// -- find in FS and read headers+body
|
||||
path := "storage/" + strconv.Itoa(id)
|
||||
path := *storage + pathsep + strconv.Itoa(id)
|
||||
|
||||
fh, err := os.Open(path + "/headers")
|
||||
fh, err := os.Open(path + pathsep + "headers")
|
||||
if err != nil {
|
||||
sendError(ctx, "Unable to read headers: ", err)
|
||||
return
|
||||
|
@ -106,7 +123,7 @@ func handler(ctx *fasthttp.RequestCtx) {
|
|||
}
|
||||
fh.Close()
|
||||
|
||||
fb, err := os.Open(path + "/body")
|
||||
fb, err := os.Open(path + pathsep + "body")
|
||||
if err != nil {
|
||||
sendError(ctx, "Unable to read body: ", err)
|
||||
return
|
||||
|
|
Loading…
Reference in a new issue