diff --git a/cmd/terse/terse.go b/cmd/terse/terse.go index de3a303..c74efa8 100644 --- a/cmd/terse/terse.go +++ b/cmd/terse/terse.go @@ -1,6 +1,7 @@ package main import ( + "fmt" "log" "github.com/brettlangdon/go-arg" @@ -10,17 +11,24 @@ import ( var args struct { MaxEntries int `arg:"-m,--max,help:max number of links to keep [default: 1000]"` Bind string `arg:"-b,--bind,help:\"[host]:\" to bind the server to [default: 127.0.0.1:5892]"` + ServerURL string `arg:"-s,--server,help:base server url to generate links as (e.g. \"https://short.domain.com\") [default: \"http://\"]"` } func main() { // Setup default args args.MaxEntries = 1000 args.Bind = "127.0.0.1:5892" - // Parse args from CLI arg.MustParse(&args) + if args.ServerURL == "" { + args.ServerURL = fmt.Sprintf("http://%s", args.Bind) + } + // Start the server - server := terse.NewServer(args.Bind, args.MaxEntries) - log.Fatal(server.ListenAndServe()) + server, err := terse.NewServer(args.Bind, args.MaxEntries, args.ServerURL) + if err == nil { + err = server.ListenAndServe() + } + log.Fatal(err) } diff --git a/server.go b/server.go index 22529d1..322c537 100644 --- a/server.go +++ b/server.go @@ -11,7 +11,8 @@ import ( ) type Handler struct { - cache *lru.Cache + cache *lru.Cache + serverURL *url.URL } func (handler *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -61,19 +62,21 @@ func (handler *Handler) HandlePost(w http.ResponseWriter, r *http.Request) { handler.cache.Add(code, cleanUrl) // Generate response url - codeUrl := &url.URL{ - Scheme: "https", - Host: r.Host, - Path: "/" + code, - } + codeUrl := handler.serverURL + codeUrl.Path = "/" + code fmt.Fprintf(w, codeUrl.String()) } -func NewServer(bind string, maxEntries int) *http.Server { +func NewServer(bind string, maxEntries int, serverURL string) (*http.Server, error) { + parsedURL, err := url.ParseRequestURI(serverURL) + if err != nil { + return nil, err + } return &http.Server{ Addr: bind, Handler: &Handler{ - cache: lru.New(maxEntries), + cache: lru.New(maxEntries), + serverURL: parsedURL, }, - } + }, nil }