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