Lay out groundwork for multiple iterations/steps

Add checking for bad solutions.
Do http requests with contexts
Minor fixes with managing steps counter.

README.md

@@ -2,7 +2,7 @@
 
 A fast [Tartarus](https://usips.org/products/tartarus/) solver lib.
 
-For technical details, you can try asking Josh (not me) or reverse engineer the *wonderful* solver solutions offered on Tartarus' no-js page. Good luck.
+For technical details, you can try asking Josh (not me) or reverse engineer the solver solutions offered on Tartarus' no-js page. Good luck.
 
 <hr>
 

cerberus.go

@@ -13,63 +13,35 @@ import (
 type Challenge struct {
 	Salt  string // Challenge salt from server.
 	Diff  uint32 // Difficulty level.
-	Steps uint32 // Not 100% sure how this is used yet.
+	Steps int8   // Each step consists of a Challenge and a Solution. More than 1 may be required.
+	// Stored as a signed int to get past underflow issues later on.
+
 	host *url.URL
 }
 
+// Convenience wrapper, equivalent to Solve(ctx, c).
+func (c Challenge) Solve(ctx context.Context) (Solution, error) {
+	return Solve(ctx, c)
+}
+
 type Solution struct {
-	Salt     string
-	Nonce    uint32
-	Redirect string // To be set manually by caller.
+	Hash     []byte // Not required in POST, but provided for reference.
+	Salt     string // Challenge salt.
+	Nonce    uint32 // Solution nonce. This is the "answer" to the problem.
+	Redirect string // Relative path to redirect to after Solution is accepted.
+
+	Steps int8 // Steps, as described in Challenge.
+
+	// TODO: Maybe make the redirect shit auto, idk.
 
-	Hash []byte
 	host *url.URL
 }
 
-// Brute force nonces until a valid solution is found.
-func genHashes(ctx context.Context, c Challenge) <-chan Solution {
-	var (
-		out      = make(chan Solution, 1)
-		sha      = sha256.New()
-		nonce    = rand.Uint32()
-	)
-
-	go func() {
-		defer close(out)
-
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			default:
-			}
-
-			sha.Write(fmt.Append(nil, c.Salt, nonce))
-
-			sol := Solution{
-				Salt:     c.Salt,
-				Nonce:    nonce,
-				Redirect: "/", // Use sensible default for placeholder.
-				Hash:     sha.Sum(nil),
-			}
-			// Ensure we don't hang if out channel is full on ctx close.
-			select {
-			case <-ctx.Done():
-				return
-			case out <- sol:
-			}
-
-			// Reset hasher input for next iteration.
-			sha.Reset()
-			nonce++
-		}
-	}()
-
-	return out
-}
-
 // Given difficulty is measured in number of leading 0 bits.
 func checkZeros(diff uint32, hash []byte) bool {
+	// I am using this big ugly check to avoid needing to hardcode the max difficulty (32, by the looks of it).
+	// Otherwise, it's a simple < comparison to (1 << (32 - diff))
+
 	var (
 		rem    = diff % 8         // Remainder after dividing diff (given in bits) to bytes.
 		nbytes = (diff - rem) / 8 // Amount of 0x0 bytes we can divide difficulty bits into.
@@ -77,9 +49,8 @@ func checkZeros(diff uint32, hash []byte) bool {
 		mask uint8 // Mask to check remaining bits.
 	)
 
-	lh := uint32(len(hash))
 	// Check bounds for the loops found below.
-	if lh < nbytes || (rem > 0 && lh < nbytes+1) {
+	if lh := uint32(len(hash)); lh < nbytes || (rem > 0 && lh < nbytes+1) {
 		return false
 	}
 
@@ -106,16 +77,57 @@ func checkZeros(diff uint32, hash []byte) bool {
 	return hash[nbytes]&mask == 0x0
 }
 
+// Brute force nonces until a valid solution is found.
+func genHashes(ctx context.Context, c Challenge) <-chan Solution {
+	var (
+		out   = make(chan Solution, 1) // Probably unnecessary buffering for questionable efficiency.
+		sha   = sha256.New()
+		nonce = rand.Uint32()
+	)
+
+	go func() {
+		defer close(out)
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
+
+			sha.Write(fmt.Append(nil, c.Salt, nonce))
+
+			sol := Solution{
+				Hash:     sha.Sum(nil),
+				Salt:     c.Salt,
+				Nonce:    nonce,
+				Redirect: "/", // Use sensible default for placeholder, for now.
+				Steps:    c.Steps - 1,
+			}
+			// Ensure we don't hang if out channel is full on ctx close.
+			select {
+			case <-ctx.Done():
+				return
+			case out <- sol:
+			}
+
+			// Reset hasher input for next iteration.
+			sha.Reset()
+			nonce++
+		}
+	}()
+
+	return out
+}
+
 // Solve Challenge c. Returns Solution that can be submitted.
 func Solve(ctx context.Context, c Challenge) (Solution, error) {
 	sol := make(chan Solution, 1)
 
-	go func() {
-		// A reasonable hardware-based job limiter.
-		// Use 1 worker per thread at most.
-		threads := runtime.NumCPU()
-		for i := 0; i < threads; i++ {
-			go func() {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	worker := func() {
 		hf := genHashes(ctx, c)
 		// Loop until answer has been found.
 		// Should break when hash worker terminates.
@@ -127,7 +139,14 @@ func Solve(ctx context.Context, c Challenge) (Solution, error) {
 				sol <- h
 			}
 		}
-			}()
+	}
+
+	go func() {
+		// A reasonable hardware-based job limiter.
+		// Use 1 worker per thread at most.
+		threads := runtime.NumCPU()
+		for range threads {
+			go worker()
 		}
 	}()
 

cerberus_test.go

@@ -32,7 +32,7 @@ func solveTest(ctx context.Context, hc http.Client, host string) error {
 	}
 
 	log.Printf("Fetching new %s challenge...", connType)
-	c, err := NewChallenge(hc, host)
+	c, err := NewChallenge(ctx, hc, host)
 	if err != nil {
 		return err
 	}
@@ -42,13 +42,15 @@ func solveTest(ctx context.Context, hc http.Client, host string) error {
 	if err != nil {
 		return err
 	}
-	log.Printf("Solution hash: %x, nonce: %d\n", s.Hash, s.Nonce)
+	log.Printf("Solution hash: %x, nonce: %d, remaining steps: %d\n", s.Hash, s.Nonce, s.Steps)
 
-	a, err := Submit(hc, s)
+	resp, err := Submit(ctx, hc, s, "")
 	if err != nil {
 		return err
 	}
-	log.Printf("Response: %s\n\n", a)
+	defer resp.Body.Close()
+
+	log.Printf("Response: %s\n\n", resp.Header.Get("Set-Cookie"))
 
 	return nil
 }
@@ -65,17 +67,13 @@ func newProxyTransport() *http.Transport {
 func TestSubmit(t *testing.T) {
 	ctx := t.Context()
 
-	hc := http.Client{}
-
-	err := solveTest(ctx, hc, _TEST_HOST)
+	err := solveTest(ctx, http.Client{}, _TEST_HOST)
 	if err != nil {
 		t.Error(err)
 	}
 
-	hc.Transport = newProxyTransport()
-
 	var dnsErr *net.DNSError
-	err = solveTest(ctx, hc, _TEST_ONION)
+	err = solveTest(ctx, http.Client{Transport: newProxyTransport()}, _TEST_ONION)
 	if err != nil {
 		if errors.As(err, &dnsErr) {
 			log.Println("Unable to resolve .onion domain. Make sure ALL_PROXY is set and tor is running.")

http.go

@@ -1,6 +1,7 @@
 package cerberus
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -18,7 +19,16 @@ var (
 	ErrParseFailed = errors.New("Failed to parse challenge from HTML data tags.")
 )
 
-func NewChallenge(hc http.Client, host string) (Challenge, error) {
+type ErrInvalidSolution struct {
+	s Solution
+}
+
+func (e *ErrInvalidSolution) Error() string {
+	return fmt.Sprintf("Received 400 status when submitting solution: %+v", e.s)
+}
+
+// Request new Tartarus challenge from provided host.
+func NewChallenge(ctx context.Context, hc http.Client, host string) (Challenge, error) {
 	u, err := parseHost(host)
 	if err != nil {
 		return Challenge{}, err
@@ -34,13 +44,18 @@ func NewChallenge(hc http.Client, host string) (Challenge, error) {
 		return nil
 	}
 
-	resp, err := hc.Get(u.String())
+	req, err := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
+	if err != nil {
+		return Challenge{}, err
+	}
+
+	resp, err := hc.Do(req)
 	if err != nil {
 		return Challenge{}, err
 	}
 	defer resp.Body.Close()
 
-	// Check for 203 status
+	// Check for 203 status. A 203 indicates a redirect to a challenge page.
 	if resp.StatusCode != 203 {
 		return Challenge{}, ErrNoRedirect
 	}
@@ -55,28 +70,39 @@ func NewChallenge(hc http.Client, host string) (Challenge, error) {
 	return c, nil
 }
 
-func Submit(hc http.Client, s Solution) (string, error) {
-	resp, err := postSolution(hc, s)
-	if err != nil {
-		return "", err
-	}
-	defer resp.Body.Close()
-
-	return resp.Header.Get("Set-Cookie"), nil
+// Submit a Solution for a Challenge.
+//
+// If redirect is empty, "/" is used as a sensible default.
+// Auth cookies get set automatically in http.Client's CookieJar.
+// The *http.Response is provided to support more advanced setups. The auth token can also be found in its Body.
+func Submit(ctx context.Context, hc http.Client, s Solution, redirect string) (*http.Response, error) {
+	if redirect != "" {
+		s.Redirect = redirect
 	}
 
-func postSolution(hc http.Client, s Solution) (*http.Response, error) {
-	// Ensure the POST url parses properly before passing the string.
-	u, err := url.Parse(fmt.Sprintf("%s://%s/.ttrs/challenge", s.host.Scheme, s.host.Hostname()))
+	resp, err := postSolution(ctx, hc, s)
 	if err != nil {
 		return nil, err
 	}
 
-	return hc.PostForm(u.String(), url.Values{
-		"salt":     []string{s.Salt},
-		"redirect": []string{s.Redirect},
-		"nonce":    []string{fmt.Sprint(s.Nonce)},
-	})
+	if s.Steps > 0 {
+		c, err := NewChallenge(ctx, hc, s.host.String())
+		if err != nil {
+			return nil, err
+		}
+
+		// maybe useful later. idk.
+		// c.Steps = s.Steps
+
+		s, err := Solve(ctx, c)
+		if err != nil {
+			return nil, err
+		}
+
+		return Submit(ctx, hc, s, redirect)
+	}
+
+	return resp, nil
 }
 
 func parseHost(addr string) (*url.URL, error) {
@@ -85,12 +111,35 @@ func parseHost(addr string) (*url.URL, error) {
 		addr = "https://" + addr
 	}
 
-	u, err := url.Parse(addr)
+	return url.Parse(addr)
+}
+
+func postSolution(ctx context.Context, hc http.Client, s Solution) (*http.Response, error) {
+	// Ensure the POST url parses properly before passing the string.
+	u, err := url.Parse(fmt.Sprintf("%s://%s/.ttrs/challenge", s.host.Scheme, s.host.Hostname()))
 	if err != nil {
 		return nil, err
 	}
 
-	return u, nil
+	reqBody := strings.NewReader(fmt.Sprintf(`salt=%s&redirect=%s&nonce=%d`, s.Salt, url.PathEscape(s.Redirect), s.Nonce))
+	req, err := http.NewRequestWithContext(ctx, "POST", u.String(), reqBody)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := hc.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO: Additionally verify failure from response JSON. Maybe include resp body in err type.
+	// Rejected solution response: status=400 body={"success":false,"reason":"invalid_solution","action":"retry"}
+	if resp.StatusCode == 400 {
+		defer resp.Body.Close()
+		return nil, &ErrInvalidSolution{s}
+	}
+
+	return resp, nil
 }
 
 func parseTags(r io.Reader) (Challenge, error) {
@@ -115,7 +164,7 @@ func parseTags(r io.Reader) (Challenge, error) {
 					if err != nil {
 						return c, ErrParseFailed
 					}
-					c.Steps = uint32(steps)
+					c.Steps = int8(steps)
 				}
 			}
 		}