1package prover
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"fmt"
  7	"math/big"
  8	"net/http"
  9	"strings"
 10	"time"
 11
 12	"github.com/dghubble/sling"
 13	"github.com/hermeznetwork/hermez-node/common"
 14	"github.com/hermeznetwork/tracerr"
 15)
 16
 17// Proof TBD this type will be received from the proof server
 18type Proof struct {
 19	PiA      [3]*big.Int    `json:"pi_a"`
 20	PiB      [3][2]*big.Int `json:"pi_b"`
 21	PiC      [3]*big.Int    `json:"pi_c"`
 22	Protocol string         `json:"protocol"`
 23}
 24
 25type bigInt big.Int
 26
 27func (b *bigInt) UnmarshalText(text []byte) error {
 28	_, ok := (*big.Int)(b).SetString(string(text), 10)
 29	if !ok {
 30		return tracerr.Wrap(fmt.Errorf("invalid big int: \"%v\"", string(text)))
 31	}
 32	return nil
 33}
 34
 35// UnmarshalJSON unmarshals the proof from a JSON encoded proof with the big
 36// ints as strings
 37func (p *Proof) UnmarshalJSON(data []byte) error {
 38	proof := struct {
 39		PiA      [3]*bigInt    `json:"pi_a"`
 40		PiB      [3][2]*bigInt `json:"pi_b"`
 41		PiC      [3]*bigInt    `json:"pi_c"`
 42		Protocol string        `json:"protocol"`
 43	}{}
 44	if err := json.Unmarshal(data, &proof); err != nil {
 45		return tracerr.Wrap(err)
 46	}
 47	p.PiA[0] = (*big.Int)(proof.PiA[0])
 48	p.PiA[1] = (*big.Int)(proof.PiA[1])
 49	p.PiA[2] = (*big.Int)(proof.PiA[2])
 50	if p.PiA[2].Int64() != 1 {
 51		return tracerr.Wrap(fmt.Errorf("Expected PiA[2] == 1, but got %v", p.PiA[2]))
 52	}
 53	p.PiB[0][0] = (*big.Int)(proof.PiB[0][0])
 54	p.PiB[0][1] = (*big.Int)(proof.PiB[0][1])
 55	p.PiB[1][0] = (*big.Int)(proof.PiB[1][0])
 56	p.PiB[1][1] = (*big.Int)(proof.PiB[1][1])
 57	p.PiB[2][0] = (*big.Int)(proof.PiB[2][0])
 58	p.PiB[2][1] = (*big.Int)(proof.PiB[2][1])
 59	if p.PiB[2][0].Int64() != 1 || p.PiB[2][1].Int64() != 0 {
 60		return tracerr.Wrap(fmt.Errorf("Expected PiB[2] == [1, 0], but got %v", p.PiB[2]))
 61	}
 62	p.PiC[0] = (*big.Int)(proof.PiC[0])
 63	p.PiC[1] = (*big.Int)(proof.PiC[1])
 64	p.PiC[2] = (*big.Int)(proof.PiC[2])
 65	if p.PiC[2].Int64() != 1 {
 66		return tracerr.Wrap(fmt.Errorf("Expected PiC[2] == 1, but got %v", p.PiC[2]))
 67	}
 68	p.Protocol = proof.Protocol
 69	return nil
 70}
 71
 72// PublicInputs are the public inputs of the proof
 73type PublicInputs []*big.Int
 74
 75// UnmarshalJSON unmarshals the JSON into the public inputs where the bigInts
 76// are in decimal as quoted strings
 77func (p *PublicInputs) UnmarshalJSON(data []byte) error {
 78	pubInputs := []*bigInt{}
 79	if err := json.Unmarshal(data, &pubInputs); err != nil {
 80		return tracerr.Wrap(err)
 81	}
 82	*p = make([]*big.Int, len(pubInputs))
 83	for i, v := range pubInputs {
 84		([]*big.Int)(*p)[i] = (*big.Int)(v)
 85	}
 86	return nil
 87}
 88
 89// Client is the interface to a ServerProof that calculates zk proofs
 90type Client interface {
 91	// Non-blocking
 92	CalculateProof(ctx context.Context, zkInputs *common.ZKInputs) error
 93	// Blocking.  Returns the Proof and Public Data (public inputs)
 94	GetProof(ctx context.Context) (*Proof, []*big.Int, error)
 95	// Non-Blocking
 96	Cancel(ctx context.Context) error
 97	// Blocking
 98	WaitReady(ctx context.Context) error
 99}
100
101// StatusCode is the status string of the ProofServer
102type StatusCode string
103
104const (
105	// StatusCodeAborted means prover is ready to take new proof. Previous
106	// proof was aborted.
107	StatusCodeAborted StatusCode = "aborted"
108	// StatusCodeBusy means prover is busy computing proof.
109	StatusCodeBusy StatusCode = "busy"
110	// StatusCodeFailed means prover is ready to take new proof. Previous
111	// proof failed
112	StatusCodeFailed StatusCode = "failed"
113	// StatusCodeSuccess means prover is ready to take new proof. Previous
114	// proof succeeded
115	StatusCodeSuccess StatusCode = "success"
116	// StatusCodeUnverified means prover is ready to take new proof.
117	// Previous proof was unverified
118	StatusCodeUnverified StatusCode = "unverified"
119	// StatusCodeUninitialized means prover is not initialized
120	StatusCodeUninitialized StatusCode = "uninitialized"
121	// StatusCodeUndefined means prover is in an undefined state. Most
122	// likely is booting up. Keep trying
123	StatusCodeUndefined StatusCode = "undefined"
124	// StatusCodeInitializing means prover is initializing and not ready yet
125	StatusCodeInitializing StatusCode = "initializing"
126	// StatusCodeReady means prover initialized and ready to do first proof
127	StatusCodeReady StatusCode = "ready"
128)
129
130// IsReady returns true when the prover is ready
131func (status StatusCode) IsReady() bool {
132	if status == StatusCodeAborted || status == StatusCodeFailed || status == StatusCodeSuccess ||
133		status == StatusCodeUnverified || status == StatusCodeReady {
134		return true
135	}
136	return false
137}
138
139// IsInitialized returns true when the prover is initialized
140func (status StatusCode) IsInitialized() bool {
141	if status == StatusCodeUninitialized || status == StatusCodeUndefined ||
142		status == StatusCodeInitializing {
143		return false
144	}
145	return true
146}
147
148// Status is the return struct for the status API endpoint
149type Status struct {
150	Status  StatusCode `json:"status"`
151	Proof   string     `json:"proof"`
152	PubData string     `json:"pubData"`
153}
154
155// ErrorServer is the return struct for an API error
156type ErrorServer struct {
157	Status  StatusCode `json:"status"`
158	Message string     `json:"msg"`
159}
160
161// Error message for ErrorServer
162func (e ErrorServer) Error() string {
163	return fmt.Sprintf("server proof status (%v): %v", e.Status, e.Message)
164}
165
166type apiMethod string
167
168const (
169	// GET is an HTTP GET
170	GET apiMethod = "GET"
171	// POST is an HTTP POST with maybe JSON body
172	POST apiMethod = "POST"
173)
174
175// ProofServerClient contains the data related to a ProofServerClient
176type ProofServerClient struct {
177	URL          string
178	client       *sling.Sling
179	pollInterval time.Duration
180}
181
182// NewProofServerClient creates a new ServerProof
183func NewProofServerClient(URL string, pollInterval time.Duration) *ProofServerClient {
184	if URL[len(URL)-1] != '/' {
185		URL += "/"
186	}
187	client := sling.New().Base(URL)
188	return &ProofServerClient{URL: URL, client: client, pollInterval: pollInterval}
189}
190
191func (p *ProofServerClient) apiRequest(ctx context.Context, method apiMethod, path string,
192	body interface{}, ret interface{}) error {
193	path = strings.TrimPrefix(path, "/")
194	var errSrv ErrorServer
195	var req *http.Request
196	var err error
197	switch method {
198	case GET:
199		req, err = p.client.New().Get(path).Request()
200	case POST:
201		req, err = p.client.New().Post(path).BodyJSON(body).Request()
202	default:
203		return tracerr.Wrap(fmt.Errorf("invalid http method: %v", method))
204	}
205	if err != nil {
206		return tracerr.Wrap(err)
207	}
208	res, err := p.client.Do(req.WithContext(ctx), ret, &errSrv)
209	if err != nil {
210		return tracerr.Wrap(err)
211	}
212	defer res.Body.Close() //nolint:errcheck
213	if !(200 <= res.StatusCode && res.StatusCode < 300) {
214		return tracerr.Wrap(errSrv)
215	}
216	return nil
217}
218
219func (p *ProofServerClient) apiStatus(ctx context.Context) (*Status, error) {
220	var status Status
221	return &status, tracerr.Wrap(p.apiRequest(ctx, GET, "/status", nil, &status))
222}
223
224func (p *ProofServerClient) apiCancel(ctx context.Context) error {
225	return tracerr.Wrap(p.apiRequest(ctx, POST, "/cancel", nil, nil))
226}
227
228func (p *ProofServerClient) apiInput(ctx context.Context, zkInputs *common.ZKInputs) error {
229	return tracerr.Wrap(p.apiRequest(ctx, POST, "/input", zkInputs, nil))
230}
231
232// CalculateProof sends the *common.ZKInputs to the ServerProof to compute the
233// Proof
234func (p *ProofServerClient) CalculateProof(ctx context.Context, zkInputs *common.ZKInputs) error {
235	return tracerr.Wrap(p.apiInput(ctx, zkInputs))
236}
237
238// GetProof retrieves the Proof and Public Data (public inputs) from the
239// ServerProof, blocking until the proof is ready.
240func (p *ProofServerClient) GetProof(ctx context.Context) (*Proof, []*big.Int, error) {
241	if err := p.WaitReady(ctx); err != nil {
242		return nil, nil, tracerr.Wrap(err)
243	}
244	status, err := p.apiStatus(ctx)
245	if err != nil {
246		return nil, nil, tracerr.Wrap(err)
247	}
248	if status.Status == StatusCodeSuccess {
249		var proof Proof
250		if err := json.Unmarshal([]byte(status.Proof), &proof); err != nil {
251			return nil, nil, tracerr.Wrap(err)
252		}
253		var pubInputs PublicInputs
254		if err := json.Unmarshal([]byte(status.PubData), &pubInputs); err != nil {
255			return nil, nil, tracerr.Wrap(err)
256		}
257		return &proof, pubInputs, nil
258	}
259	return nil, nil, tracerr.Wrap(fmt.Errorf("status != %v, status = %v", StatusCodeSuccess,
260		status.Status))
261}
262
263// Cancel cancels any current proof computation
264func (p *ProofServerClient) Cancel(ctx context.Context) error {
265	return tracerr.Wrap(p.apiCancel(ctx))
266}
267
268// WaitReady waits until the serverProof is ready
269func (p *ProofServerClient) WaitReady(ctx context.Context) error {
270	for {
271		status, err := p.apiStatus(ctx)
272		if err != nil {
273			return tracerr.Wrap(err)
274		}
275		if !status.Status.IsInitialized() {
276			return tracerr.Wrap(fmt.Errorf("Proof Server is not initialized"))
277		}
278		if status.Status.IsReady() {
279			return nil
280		}
281		select {
282		case <-ctx.Done():
283			return tracerr.Wrap(common.ErrDone)
284		case <-time.After(p.pollInterval):
285		}
286	}
287}
288
289// MockClient is a mock ServerProof to be used in tests.  It doesn't calculate anything
290type MockClient struct {
291	counter int64
292	Delay   time.Duration
293}
294
295// CalculateProof sends the *common.ZKInputs to the ServerProof to compute the
296// Proof
297func (p *MockClient) CalculateProof(ctx context.Context, zkInputs *common.ZKInputs) error {
298	return nil
299}
300
301// GetProof retrieves the Proof from the ServerProof
302func (p *MockClient) GetProof(ctx context.Context) (*Proof, []*big.Int, error) {
303	// Simulate a delay
304	select {
305	case <-time.After(p.Delay): //nolint:gomnd
306		i := p.counter * 100 //nolint:gomnd
307		p.counter++
308		return &Proof{
309				PiA: [3]*big.Int{
310					big.NewInt(i), big.NewInt(i + 1), big.NewInt(1), //nolint:gomnd
311				},
312				PiB: [3][2]*big.Int{
313					{big.NewInt(i + 2), big.NewInt(i + 3)}, //nolint:gomnd
314					{big.NewInt(i + 4), big.NewInt(i + 5)}, //nolint:gomnd
315					{big.NewInt(1), big.NewInt(0)},         //nolint:gomnd
316				},
317				PiC: [3]*big.Int{
318					big.NewInt(i + 6), big.NewInt(i + 7), big.NewInt(1), //nolint:gomnd
319				},
320				Protocol: "groth",
321			},
322			[]*big.Int{big.NewInt(i + 42)}, //nolint:gomnd
323			nil
324	case <-ctx.Done():
325		return nil, nil, tracerr.Wrap(common.ErrDone)
326	}
327}
328
329// Cancel cancels any current proof computation
330func (p *MockClient) Cancel(ctx context.Context) error {
331	// Simulate a delay
332	select {
333	case <-time.After(80 * time.Millisecond): //nolint:gomnd
334		return nil
335	case <-ctx.Done():
336		return tracerr.Wrap(common.ErrDone)
337	}
338}
339
340// WaitReady waits until the prover is ready
341func (p *MockClient) WaitReady(ctx context.Context) error {
342	// Simulate a delay
343	select {
344	case <-time.After(200 * time.Millisecond): //nolint:gomnd
345		return nil
346	case <-ctx.Done():
347		return tracerr.Wrap(common.ErrDone)
348	}
349}