Refactor

Author: runxiyu

bitwise.go

@@ -3,6 +3,7 @@
 
 package main
 
+// validateBitZeros checks if the first n bits of a byte slice are all zeros.
 func validateBitZeros(bs []byte, n uint) bool {
 	q := n / 8
 	r := n % 8

flags.go

@@ -11,6 +11,7 @@ var (
 	secondary  bool
 )
 
+// This init parses command line flags.
 func init() {
 	flag.UintVar(&global.NeedBits, "difficulty", 20, "leading zero bits required for the challenge")
 	flag.StringVar(&global.SourceURL, "source", "https://forge.lindenii.runxiyu.org/powxy/:/repos/powxy/", "url to the source code")

global.go

@@ -3,11 +3,16 @@
 
 package main
 
+// global is a struct that holds global information that the HTML template may
+// wish to access.
 var global = struct {
+	// Most of these fields should be filled in by flag parsing.
 	NeedBits        uint
 	NeedBitsReverse uint
 	SourceURL       string
 	Version         string
 }{
+	// The version should be replaced by the init function in version.go,
+	// if version.go is properly generated by the Makefile.
 	Version: "(no version)",
 }

identifier.go

@@ -6,11 +6,15 @@ package main
 import (
 	"crypto/hmac"
 	"crypto/sha256"
+	"crypto/subtle"
+	"encoding/base64"
 	"encoding/binary"
 	"net/http"
 	"time"
 )
 
+// makeIdentifierMAC generates an identifier that semi-uniquely identifies the client,
+// and generates a MAC for that identifier.
 func makeIdentifierMAC(request *http.Request) (identifier []byte, mac []byte) {
 	identifier = make([]byte, 0, sha256.Size)
 	mac = make([]byte, 0, sha256.Size)
@@ -37,3 +41,18 @@ func makeIdentifierMAC(request *http.Request) (identifier []byte, mac []byte) {
 
 	return
 }
+
+// validateCookie checks if the cookie is valid by comparing the base64-decoded
+// value of the cookie with an expected MAC.
+func validateCookie(cookie *http.Cookie, expectedMAC []byte) bool {
+	if cookie == nil {
+		return false
+	}
+
+	gotMAC, err := base64.StdEncoding.DecodeString(cookie.Value)
+	if err != nil {
+		return false
+	}
+
+	return subtle.ConstantTimeCompare(gotMAC, expectedMAC) == 1
+}

ip.go

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: BSD-2-Clause
+// SPDX-FileCopyrightText: Copyright (c) 2025 Runxi Yu <https://runxiyu.org>
+
+package main
+
+import (
+	"net/http"
+	"strings"
+)
+
+// getRemoteIP returns the remote IP address of the client. It respects
+// X-Forwarded-For if powxy is configured as a secondary proxy. Ports
+// are stripped.
+func getRemoteIP(request *http.Request) (remoteIP string) {
+	if secondary {
+		remoteIP, _, _ = strings.Cut(request.Header.Get("X-Forwarded-For"), ",")
+	}
+	if remoteIP == "" {
+		remoteIP = request.RemoteAddr
+		index := strings.LastIndex(remoteIP, ":")
+		if index != -1 {
+			remoteIP = remoteIP[:index]
+		}
+	}
+	return
+}

main.go

@@ -4,8 +4,6 @@
 package main
 
 import (
-	"crypto/sha256"
-	"crypto/subtle"
 	"encoding/base64"
 	"errors"
 	"log"
@@ -21,29 +19,39 @@ type tparams struct {
 
 func main() {
 	log.Fatal(http.ListenAndServe(listenAddr, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+		// Static resources for powxy itself.
 		if strings.HasPrefix(request.URL.Path, "/.powxy/") {
 			http.StripPrefix("/.powxy/", http.FileServer(http.FS(resourcesFS))).ServeHTTP(writer, request)
 			return
 		}
 
+		// We attempt to fetch the powxy cookie. Its non-existence
+		// does not matter here; if the cookie does not exist, it
+		// will be nil, so validation will simply fail and the user
+		// will be prompted to solve the PoW challenge.
 		cookie, err := request.Cookie("powxy")
-		if err != nil {
-			if !errors.Is(err, http.ErrNoCookie) {
-				log.Println("COOKIE_ERR", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
-				http.Error(writer, "error fetching cookie", http.StatusInternalServerError)
-				return
-			}
+		if err != nil && !errors.Is(err, http.ErrNoCookie) {
+			log.Println("COOKIE_ERR", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
+			http.Error(writer, "error fetching cookie", http.StatusInternalServerError)
+			return
 		}
 
+		// We generate the identifier that identifies the client,
+		// and the expected HMAC that the cookie should include.
 		identifier, expectedMAC := makeIdentifierMAC(request)
 
+		// If the cookie exists and is valid, we simply proxy the
+		// request.
 		if validateCookie(cookie, expectedMAC) {
 			log.Println("PROXY", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
 			proxyRequest(writer, request)
 			return
 		}
 
-		authPage := func(message string) {
+		// A convenience function to render the challenge page,
+		// since all parameters but the message are constant at this
+		// point.
+		challengePage := func(message string) {
 			err := tmpl.Execute(writer, tparams{
 				Identifier: base64.StdEncoding.EncodeToString(identifier),
 				Message:    message,
@@ -54,46 +62,59 @@ func main() {
 			}
 		}
 
+		// This generally shouldn't happen, at least not for web
+		// browesrs.
 		if request.ParseForm() != nil {
 			log.Println("MALFORMED", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
-			authPage("You submitted a malformed form.")
+			challengePage("You submitted a malformed form.")
 			return
 		}
 
 		formValues, ok := request.PostForm["powxy"]
 		if !ok {
+			// If there's simply no form value, the user is probably
+			// just visiting the site for the first time or with an
+			// expired cookie.
 			log.Println("CHALLENGE", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
-			authPage("")
+			challengePage("")
 			return
 		} else if len(formValues) != 1 {
+			// This should never happen, at least not for web
+			// browsers.
 			log.Println("FORM_VALUES", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
-			authPage("You submitted an invalid number of form values.")
+			challengePage("You submitted an invalid number of form values.")
 			return
 		}
 
-		nonce, err := base64.StdEncoding.DecodeString(formValues[0])
-		if err != nil {
-			log.Println("BASE64", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
-			authPage("Your submission was improperly encoded.")
+		// We validate that the length is reasonable before even
+		// decoding it with base64.
+		if len(formValues[0]) > 43 {
+			log.Println("TOO_LONG", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
+			challengePage("Your submission was too long.")
 			return
 		}
 
-		if len(nonce) > 32 {
-			log.Println("TOO_LONG", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
-			authPage("Your submission was too long.")
+		// Actually decode the base64 value.
+		nonce, err := base64.StdEncoding.DecodeString(formValues[0])
+		if err != nil {
+			log.Println("BASE64", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
+			challengePage("Your submission was improperly encoded.")
 			return
 		}
 
-		h := sha256.New()
-		h.Write(identifier)
-		h.Write(nonce)
-		ck := h.Sum(nil)
-		if !validateBitZeros(ck, global.NeedBits) {
+		// Validate the nonce.
+		if !validateNonce(identifier, nonce) {
 			log.Println("WRONG", getRemoteIP(request), request.RequestURI, request.Header.Get("User-Agent"))
-			authPage("Your submission was incorrect, or your session has expired while submitting.")
+			challengePage("Your submission was incorrect, or your session has expired while submitting.")
 			return
 		}
 
+		// Everything starting here: the nonce is valid, and we
+		// can set the cookie and redirect them. The redirection is
+		// needed as their "normal" request is most definitely
+		// different from one to expect after solving the PoW
+		// challenge.
+
 		http.SetCookie(writer, &http.Cookie{
 			Name:     "powxy",
 			Value:    base64.StdEncoding.EncodeToString(expectedMAC),
@@ -105,30 +126,3 @@ func main() {
 		http.Redirect(writer, request, "", http.StatusSeeOther)
 	})))
 }
-
-func validateCookie(cookie *http.Cookie, expectedMAC []byte) bool {
-	if cookie == nil {
-		return false
-	}
-
-	gotMAC, err := base64.StdEncoding.DecodeString(cookie.Value)
-	if err != nil {
-		return false
-	}
-
-	return subtle.ConstantTimeCompare(gotMAC, expectedMAC) == 1
-}
-
-func getRemoteIP(request *http.Request) (remoteIP string) {
-	if secondary {
-		remoteIP, _, _ = strings.Cut(request.Header.Get("X-Forwarded-For"), ",")
-	}
-	if remoteIP == "" {
-		remoteIP = request.RemoteAddr
-		index := strings.LastIndex(remoteIP, ":")
-		if index != -1 {
-			remoteIP = remoteIP[:index]
-		}
-	}
-	return
-}

privkey.go

@@ -10,10 +10,15 @@ import (
 )
 
 var (
-	privkey     = make([]byte, 32)
+	// The private key used to HMAC the challenge.
+	privkey = make([]byte, 32)
+
+	// The hash of the private key. We use this as an element of the
+	// identifier.
 	privkeyHash = make([]byte, 0, sha256.Size)
 )
 
+// This init generates a random private key and its hash.
 func init() {
 	if _, err := rand.Read(privkey); err != nil {
 		log.Fatal(err)

proxy.go

@@ -12,6 +12,8 @@ import (
 
 var reverseProxy *httputil.ReverseProxy
 
+// This init sets up the reverse proxy. Go's NewSingleHostReverseProxy is
+// sufficient for our use case.
 func init() {
 	parsedURL, err := url.Parse(destHost)
 	if err != nil {
@@ -20,6 +22,7 @@ func init() {
 	reverseProxy = httputil.NewSingleHostReverseProxy(parsedURL)
 }
 
+// proxyRequest proxies the incoming request to the destination host.
 func proxyRequest(writer http.ResponseWriter, request *http.Request) {
 	reverseProxy.ServeHTTP(writer, request)
 }

tmpl.go

@@ -13,6 +13,7 @@ import (
 var tmplString string
 var tmpl *template.Template
 
+// This init function parses the HTML template.
 func init() {
 	var err error
 	tmpl, err = template.New("powxy").Parse(tmplString)

unsafe.go

@@ -5,7 +5,7 @@ package main
 
 import "unsafe"
 
-// Converts a string to a byte slice without copying the string.
+// stringToBytes converts a string to a byte slice without copying the string.
 // Memory is borrowed from the string.
 // The resulting byte slice must not be modified in any form.
 func stringToBytes(s string) (bytes []byte) {

validate.go

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: BSD-2-Clause
+// SPDX-FileCopyrightText: Copyright (c) 2025 Runxi Yu <https://runxiyu.org>
+
+package main
+
+import (
+	"crypto/sha256"
+)
+
+// validateNonce checks if the nonce for the proof of work challenge is valid
+// for the given identifier.
+func validateNonce(identifier, nonce []byte) bool {
+	h := sha256.New()
+	h.Write(identifier)
+	h.Write(nonce)
+	ck := h.Sum(nil)
+	return validateBitZeros(ck, global.NeedBits)
+}