Guess destination from Referer header

2 files changed, 30 insertions, 2 deletions

diff --git a/handler.go b/handler.go
index edb0ff2..76ea010 100644
--- a/handler.go
+++ b/handler.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"github.com/valyala/fasthttp"
 	"net/http"
 )
@@ -23,8 +24,18 @@ func (h *RequestHandler) handleRequest(ctx *fasthttp.RequestCtx) {
 		if path.ResourcePath == "/_portgate" {
 			h.handlePortgateRequest(ctx)
 		} else {
-			// TODO: Try to grab actual destination from Referer header.
-			h.handleUnknownRequest(ctx)
+			// Try to grab actual destination from Referer header.
+			refererPath, err := ParsePathFromReferer(path, string(ctx.Request.Header.Referer()))
+			if err != nil || refererPath.DestinationIdentifier == -1 {
+				// The referer path also has no destination
+				h.handleUnknownRequest(ctx)
+			} else {
+				// We found the destination from the referer path, so we
+				// redirect the user to the Portgate URL they should've requested.
+
+				portgateUrl := fmt.Sprintf("/%d%s", refererPath.DestinationIdentifier, refererPath.ResourcePath)
+				ctx.Redirect(portgateUrl, http.StatusTemporaryRedirect)
+			}
 		}
 	} else {
 		// We were given a port, so we have to pass the request through to the destination host.
diff --git a/path.go b/path.go
index bb7154a..4b60dfe 100644
--- a/path.go
+++ b/path.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"net/url"
 	"path"
 	"strconv"
 	"strings"
@@ -53,6 +54,22 @@ func ParsePath(p string) Path {
 	}
 }
 
+// ParsePathFromReferer tries to create a Path from the Referer header of the request.
+func ParsePathFromReferer(p Path, r string) (Path, error) {
+	u, err := url.Parse(r)
+	if err != nil {
+		return Path{}, err
+	}
+
+	// p has the correct resource path but the wrong port, so we create a new path
+	// with the correct data from both.
+	rp := ParsePath(u.Path)
+	return Path{
+		DestinationIdentifier: rp.DestinationIdentifier,
+		ResourcePath:          p.ResourcePath,
+	}, nil
+}
+
 // MakeUrl creates the URL on the destination host that the user wants to access.
 func (p *Path) MakeUrl(targetHost string) string {
 	// TODO: Figure out what to do with TLS