golink: do not set HSTS headers when serving non-FQDN origins

Inspect the `Host` header to ensure that we do not return HSTS headers
for short domains which can lead to some clients pinning short domains to
endpoints with invalid certificates.

Signed-off-by: Patrick O'Doherty <[email protected]>

Author: patrickod

golink.go

@@ -334,10 +334,17 @@ func redirectHandler(hostname string) http.Handler {
 }
 
 // HSTS wraps the provided handler and sets Strict-Transport-Security header on
-// all responses.
+// responses. It inspects the Host header to ensure we do not specify HSTS
+// response on non fully qualified domain name origins.
 func HSTS(h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Strict-Transport-Security", "max-age=31536000")
+		host, found := r.Header["Host"]
+		if found {
+			host := host[0]
+			if strings.Contains(host, ".") {
+				w.Header().Set("Strict-Transport-Security", "max-age=31536000")
+			}
+		}
 		h.ServeHTTP(w, r)
 	})
 }

golink_test.go

@@ -524,3 +524,41 @@ func TestResolveLink(t *testing.T) {
 		})
 	}
 }
+
+func TestNoHSTSShortDomain(t *testing.T) {
+	var err error
+	db, err = NewSQLiteDB(":memory:")
+	if err != nil {
+		t.Fatal(err)
+	}
+	db.Save(&Link{Short: "foobar", Long: "http://foobar/"})
+
+	tests := []struct {
+		host       string
+		expectHsts bool
+	}{
+		{
+			host:       "go",
+			expectHsts: false,
+		},
+		{
+			host:       "go.prawn-universe.ts.net",
+			expectHsts: true,
+		},
+	}
+	for _, tt := range tests {
+		name := "HSTS: " + tt.host
+		t.Run(name, func(t *testing.T) {
+			r := httptest.NewRequest("GET", "/foobar", nil)
+			r.Header.Add("Host", tt.host)
+
+			w := httptest.NewRecorder()
+			HSTS(serveHandler()).ServeHTTP(w, r)
+
+			_, found := w.Header()["Strict-Transport-Security"]
+			if found != tt.expectHsts {
+				t.Errorf("HSTS expectation: domain %s want: %t got: %t", tt.host, tt.expectHsts, found)
+			}
+		})
+	}
+}