arnaucube
arnaucube/hermez-nodePublic Self-hosted repository
hermez-node/test/proofserver/proofserver.go
test/proofserver/proofserver.goRaw
  1package proofserver
  2
  3import (
  4	"context"
  5	"fmt"
  6	"io/ioutil"
  7	"net/http"
  8	"sync"
  9	"time"
 10
 11	"github.com/gin-contrib/cors"
 12	"github.com/gin-gonic/gin"
 13	"github.com/hermeznetwork/hermez-node/log"
 14	"github.com/hermeznetwork/hermez-node/prover"
 15	"github.com/hermeznetwork/tracerr"
 16)
 17
 18type msg struct {
 19	value string
 20	ackCh chan bool
 21}
 22
 23func newMsg(value string) msg {
 24	return msg{
 25		value: value,
 26		ackCh: make(chan bool),
 27	}
 28}
 29
 30// Mock proof server
 31type Mock struct {
 32	addr   string
 33	status prover.StatusCode
 34	sync.RWMutex
 35	proof           string
 36	pubData         string
 37	counter         int
 38	msgCh           chan msg
 39	wg              sync.WaitGroup
 40	provingDuration time.Duration
 41}
 42
 43// NewMock creates a new mock server
 44func NewMock(addr string, provingDuration time.Duration) *Mock {
 45	return &Mock{
 46		addr:            addr,
 47		status:          prover.StatusCodeReady,
 48		proof:           "",
 49		pubData:         "",
 50		counter:         0,
 51		msgCh:           make(chan msg),
 52		provingDuration: provingDuration,
 53	}
 54}
 55
 56func (s *Mock) err(c *gin.Context, err error) {
 57	c.JSON(http.StatusInternalServerError, prover.ErrorServer{
 58		Status:  "error",
 59		Message: err.Error(),
 60	})
 61}
 62
 63func (s *Mock) handleCancel(c *gin.Context) {
 64	msg := newMsg("cancel")
 65	s.msgCh <- msg
 66	<-msg.ackCh
 67	c.JSON(http.StatusOK, "OK")
 68}
 69
 70func (s *Mock) handleStatus(c *gin.Context) {
 71	s.RLock()
 72	c.JSON(http.StatusOK, prover.Status{
 73		Status:  s.status,
 74		Proof:   s.proof,
 75		PubData: s.pubData,
 76	})
 77	s.RUnlock()
 78}
 79
 80func (s *Mock) handleInput(c *gin.Context) {
 81	s.RLock()
 82	if !s.status.IsReady() {
 83		s.err(c, fmt.Errorf("not ready"))
 84		s.RUnlock()
 85		return
 86	}
 87	s.RUnlock()
 88	_, err := ioutil.ReadAll(c.Request.Body)
 89	if err != nil {
 90		s.err(c, err)
 91		return
 92	}
 93	msg := newMsg("prove")
 94	s.msgCh <- msg
 95	<-msg.ackCh
 96	c.JSON(http.StatusOK, "OK")
 97}
 98
 99const longWaitDuration = 999 * time.Hour
100
101// const provingDuration = 2 * time.Second
102
103func (s *Mock) runProver(ctx context.Context) {
104	waitDuration := longWaitDuration
105	for {
106		select {
107		case <-ctx.Done():
108			return
109		case msg := <-s.msgCh:
110			switch msg.value {
111			case "cancel":
112				waitDuration = longWaitDuration
113				s.Lock()
114				if !s.status.IsReady() {
115					s.status = prover.StatusCodeAborted
116				}
117				s.Unlock()
118			case "prove":
119				waitDuration = s.provingDuration
120				s.Lock()
121				s.status = prover.StatusCodeBusy
122				s.Unlock()
123			}
124			msg.ackCh <- true
125		case <-time.After(waitDuration):
126			waitDuration = longWaitDuration
127			s.Lock()
128			if s.status != prover.StatusCodeBusy {
129				s.Unlock()
130				continue
131			}
132			i := s.counter * 100 //nolint:gomnd
133			s.counter++
134			// Mock data
135			s.proof = fmt.Sprintf(`{
136				"pi_a": ["%v", "%v"],
137				"pi_b": [["%v", "%v"],["%v", "%v"],["%v", "%v"]],
138				"pi_c": ["%v", "%v"],
139				"protocol": "groth16"
140			}`, i, i+1, i+2, i+3, i+4, i+5, i+6, i+7, i+8, i+9) //nolint:gomnd
141			s.pubData = fmt.Sprintf(`[
142				"%v"
143			]`, i+42) //nolint:gomnd
144			s.status = prover.StatusCodeSuccess
145			s.Unlock()
146		}
147	}
148}
149
150// Run the mock server.  Use ctx to stop it via cancel
151func (s *Mock) Run(ctx context.Context) error {
152	api := gin.Default()
153	api.Use(cors.Default())
154
155	apiGroup := api.Group("/api")
156	apiGroup.GET("/status", s.handleStatus)
157	apiGroup.POST("/input", s.handleInput)
158	apiGroup.POST("/cancel", s.handleCancel)
159
160	debugAPIServer := &http.Server{
161		Addr:    s.addr,
162		Handler: api,
163		// Use some hardcoded numberes that are suitable for testing
164		ReadTimeout:    30 * time.Second, //nolint:gomnd
165		WriteTimeout:   30 * time.Second, //nolint:gomnd
166		MaxHeaderBytes: 1 << 20,          //nolint:gomnd
167	}
168	go func() {
169		log.Infof("prover.MockServer is ready at %v", s.addr)
170		if err := debugAPIServer.ListenAndServe(); err != nil && tracerr.Unwrap(err) != http.ErrServerClosed {
171			log.Fatalf("Listen: %s\n", err)
172		}
173	}()
174	s.wg.Add(1)
175	go func() {
176		s.runProver(ctx)
177		s.wg.Done()
178	}()
179
180	<-ctx.Done()
181	log.Info("Stopping prover.MockServer...")
182
183	s.wg.Wait()
184	ctxTimeout, cancel := context.WithTimeout(context.Background(), 10*time.Second) //nolint:gomnd
185	defer cancel()
186	if err := debugAPIServer.Shutdown(ctxTimeout); err != nil {
187		return tracerr.Wrap(err)
188	}
189	log.Info("prover.MockServer done")
190	return nil
191}