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