[chore] more NoLLaMas proof-of-work tweaking (#4096)

- replaces the sha256 calculation with an alternative implementation that seems to use more uniform time-taken across different platforms
- goes back to the simpler difficulty calculation without a "partial" difficulty level

Reviewed-on: https://codeberg.org/superseriousbusiness/gotosocial/pulls/4096
Co-authored-by: kim <grufwub@gmail.com>
Co-committed-by: kim <grufwub@gmail.com>

2 files changed, 26 insertions, 74 deletions

diff --git a/internal/middleware/nollamas.go b/internal/middleware/nollamas.go
index e5be014f5..eedf3b9b7 100644
--- a/internal/middleware/nollamas.go
+++ b/internal/middleware/nollamas.go
@@ -72,8 +72,7 @@ func NoLLaMas(
 	var nollamas nollamas
 	nollamas.seed = seed
 	nollamas.ttl = time.Hour
-	nollamas.diff1 = 4
-	nollamas.diff2 = '4'
+	nollamas.diff = 4
 	nollamas.getInstanceV1 = getInstanceV1
 	nollamas.policy = cookiePolicy
 	return nollamas.Serve
@@ -101,16 +100,9 @@ type nollamas struct {
 	ttl time.Duration
 
 	// algorithm difficulty knobs.
-	// diff1 determines the number of
-	// leading zeroes required, while
-	// diff2 checks the next byte at
-	// index is less than it.
-	//
-	// e.g. you look for say:
-	// - b[0:3] must be '0'
-	// - b[4] can be < '5'
-	diff1 uint8
-	diff2 uint8
+	// diff determines the number
+	// of leading zeroes required.
+	diff uint8
 
 	// extra fields required for
 	// our template rendering.
@@ -187,6 +179,12 @@ func (m *nollamas) Serve(c *gin.Context) {
 		return
 	}
 
+	// From here-on out, all
+	// possibilities are handled
+	// by us. Prevent further http
+	// handlers from being called.
+	c.Abort()
+
 	// Prepare new log entry.
 	l := log.WithContext(ctx).
 		WithField("userAgent", userAgent).
@@ -225,10 +223,6 @@ func (m *nollamas) Serve(c *gin.Context) {
 
 	l.Infof("challenge passed: %s", nonce)
 
-	// Don't pass to further
-	// handlers, we'll redirect.
-	c.Abort()
-
 	// Drop solution query and encode.
 	query.Del("nollamas_solution")
 	c.Request.URL.RawQuery = query.Encode()
@@ -240,11 +234,6 @@ func (m *nollamas) Serve(c *gin.Context) {
 }
 
 func (m *nollamas) renderChallenge(c *gin.Context, challenge string) {
-	// Don't pass to further
-	// handlers, they only get
-	// our challenge page.
-	c.Abort()
-
 	// Fetch current instance information for templating vars.
 	instance, errWithCode := m.getInstanceV1(c.Request.Context())
 	if errWithCode != nil {
@@ -263,12 +252,8 @@ func (m *nollamas) renderChallenge(c *gin.Context, challenge string) {
 			"/assets/Fork-Awesome/css/fork-awesome.min.css",
 		},
 		Extra: map[string]any{
-			"challenge":   challenge,
-			"difficulty1": m.diff1,
-
-			// must be a str otherwise template
-			// renders uint8 as int, not char
-			"difficulty2": hexStrs[m.diff2],
+			"challenge":  challenge,
+			"difficulty": m.diff,
 		},
 		Javascript: []apiutil.JavascriptEntry{
 			{
@@ -289,8 +274,7 @@ func (m *nollamas) token(hash *hashWithBufs, userAgent, clientIP string) string
 	// Include difficulty level in
 	// hash input data so if config
 	// changes then token invalidates.
-	hash.hash.Write([]byte{m.diff1})
-	hash.hash.Write([]byte{m.diff2})
+	hash.hash.Write([]byte{m.diff})
 
 	// Also seed the generated input with
 	// current time rounded to TTL, so our
@@ -326,40 +310,18 @@ func (m *nollamas) checkChallenge(hash *hashWithBufs, challenge, nonce string) b
 	hex.Encode(hash.ebuf, hash.hbuf)
 	solution := hash.ebuf
 
-	// Compiler bound-check-elimination hint.
-	if len(solution) < int(m.diff1+1) {
+	// Compiler bound-check hint.
+	if len(solution) < int(m.diff) {
 		panic(gtserror.New("BCE"))
 	}
 
 	// Check that the first 'diff'
 	// many chars are indeed zeroes.
-	for i := range m.diff1 {
+	for i := range m.diff {
 		if solution[i] != '0' {
 			return false
 		}
 	}
 
-	// Check that next char is < 'diff2'.
-	return solution[m.diff1] < m.diff2
-}
-
-// hexStrs is a quick lookup of ASCII hex
-// bytes to their string equivalent.
-var hexStrs = [...]string{
-	'0': "0",
-	'1': "1",
-	'2': "2",
-	'3': "3",
-	'4': "4",
-	'5': "5",
-	'6': "6",
-	'7': "7",
-	'8': "8",
-	'9': "9",
-	'a': "a",
-	'b': "b",
-	'c': "c",
-	'd': "d",
-	'e': "e",
-	'f': "f",
+	return true
 }
diff --git a/internal/middleware/nollamas_test.go b/internal/middleware/nollamas_test.go
index d6fdb5ff6..37b66e5e9 100644
--- a/internal/middleware/nollamas_test.go
+++ b/internal/middleware/nollamas_test.go
@@ -96,8 +96,7 @@ func testNoLLaMasMiddleware(t *testing.T, e *gin.Engine, userAgent string) {
 	}
 
 	var challenge string
-	var diff1 uint64
-	var diff2 uint8
+	var difficulty uint64
 
 	// Parse output body and find the challenge / difficulty.
 	for _, line := range strings.Split(string(b), "\n") {
@@ -107,22 +106,17 @@ func testNoLLaMasMiddleware(t *testing.T, e *gin.Engine, userAgent string) {
 			line = line[25:]
 			line = line[:len(line)-1]
 			challenge = line
-		case strings.HasPrefix(line, "data-nollamas-difficulty1=\""):
-			line = line[27:]
+		case strings.HasPrefix(line, "data-nollamas-difficulty=\""):
+			line = line[26:]
 			line = line[:len(line)-1]
 			var err error
-			diff1, err = strconv.ParseUint(line, 10, 8)
+			difficulty, err = strconv.ParseUint(line, 10, 8)
 			assert.NoError(t, err)
-		case strings.HasPrefix(line, "data-nollamas-difficulty2=\""):
-			line = line[27:]
-			line = line[:len(line)-1]
-			diff2 = line[0]
 		}
 	}
 
 	// Ensure valid posed challenge.
-	assert.NotZero(t, diff1)
-	assert.NotZero(t, diff2)
+	assert.NotZero(t, difficulty)
 	assert.NotEmpty(t, challenge)
 
 	// Prepare a test request for gin engine.
@@ -131,12 +125,11 @@ func testNoLLaMasMiddleware(t *testing.T, e *gin.Engine, userAgent string) {
 	rw = httptest.NewRecorder()
 
 	// Now compute and set solution query paramater.
-	solution := computeSolution(challenge, diff1, diff2)
+	solution := computeSolution(challenge, difficulty)
 	r.URL.RawQuery = "nollamas_solution=" + solution
 
 	t.Logf("challenge=%s", challenge)
-	t.Logf("diff1=%d", diff1)
-	t.Logf("diff2='%c'", diff2)
+	t.Logf("difficulty=%d", difficulty)
 	t.Logf("solution=%s", solution)
 
 	// Pass req through
@@ -159,21 +152,18 @@ func testNoLLaMasMiddleware(t *testing.T, e *gin.Engine, userAgent string) {
 }
 
 // computeSolution does the functional equivalent of our nollamas workerTask.js.
-func computeSolution(challenge string, diff1 uint64, diff2 uint8) string {
+func computeSolution(challenge string, diff uint64) string {
 outer:
 	for i := 0; ; i++ {
 		solution := strconv.Itoa(i)
 		combined := challenge + solution
 		hash := sha256.Sum256(byteutil.S2B(combined))
 		encoded := hex.EncodeToString(hash[:])
-		for i := range diff1 {
+		for i := range diff {
 			if encoded[i] != '0' {
 				continue outer
 			}
 		}
-		if encoded[diff1] >= diff2 {
-			continue outer
-		}
 		return solution
 	}
 }