Merge pull request #15 from Shopify/v3-client

V3 client
Shopify · May 20, 2022 · b0753dc · b0753dc
2 parents b4d994f + c12b43d
commit b0753dc
Show file tree

Hide file tree

Showing 8 changed files with 201 additions and 179 deletions.
diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
@@ -12,9 +12,12 @@ jobs:
     name: golang lint
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/setup-go@v3
+        with:
+          go-version: 1.16
+      - uses: actions/checkout@v3
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v2
+        uses: golangci/golangci-lint-action@v3
         with:
           version: v1.39
 

diff --git a/client.go b/client.go
@@ -4,10 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/go-resty/resty/v2"
+	"io"
 	"log"
 	"net/http"
-
-	"github.com/go-resty/resty/v2"
+	"strconv"
 )
 
 const (
@@ -17,15 +18,19 @@ const (
 	learnSpamEndpoint = "learnspam"
 	learnHamEndpoint  = "learnham"
 	pingEndpoint      = "ping"
+
+	QueueID = "Queue-Id"
+	Flag    = "Flag"
+	Weight  = "Weight"
 )
 
 // Client is a rspamd HTTP client.
 type Client interface {
-	Check(context.Context, *Email) (*CheckResponse, error)
-	LearnSpam(context.Context, *Email) (*LearnResponse, error)
-	LearnHam(context.Context, *Email) (*LearnResponse, error)
-	FuzzyAdd(context.Context, *Email) (*LearnResponse, error)
-	FuzzyDel(context.Context, *Email) (*LearnResponse, error)
+	Check(context.Context, *CheckRequest) (*CheckResponse, error)
+	LearnSpam(context.Context, *LearnRequest) (*LearnResponse, error)
+	LearnHam(context.Context, *LearnRequest) (*LearnResponse, error)
+	FuzzyAdd(context.Context, *FuzzyRequest) (*FuzzyResponse, error)
+	FuzzyDel(context.Context, *FuzzyRequest) (*FuzzyResponse, error)
 	Ping(context.Context) (PingResponse, error)
 }
 
@@ -35,18 +40,53 @@ type client struct {
 
 var _ Client = &client{}
 
+// CheckRequest encapsulates the request of Check.
+type CheckRequest struct {
+	Message io.Reader
+	Header  http.Header
+}
+
+// SymbolData encapsulates the data returned for each symbol from Check.
+type SymbolData struct {
+	Name        string  `json:"name"`
+	Score       float64 `json:"score"`
+	MetricScore float64 `json:"metric_score"`
+	Description string  `json:"description"`
+}
+
 // CheckResponse encapsulates the response of Check.
 type CheckResponse struct {
 	Score     float64               `json:"score"`
+	Action    string                `json:"action"`
 	MessageID string                `json:"message-id"`
 	Symbols   map[string]SymbolData `json:"symbols"`
 }
 
-// LearnResponse encapsulates the response of LearnSpam, LearnHam, FuzzyAdd, FuzzyDel.
+// LearnRequest encapsulates the request of LearnSpam, LearnHam.
+type LearnRequest struct {
+	Message io.Reader
+	Header  http.Header
+}
+
+// LearnResponse encapsulates the response of LearnSpam, LearnHam.
 type LearnResponse struct {
 	Success bool `json:"success"`
 }
 
+// FuzzyRequest encapsulates the request of FuzzyAdd, FuzzyDel.
+type FuzzyRequest struct {
+	Message io.Reader
+	Flag    int
+	Weight  int
+	Header  http.Header
+}
+
+// FuzzyResponse encapsulates the response of FuzzyAdd, FuzzyDel.
+type FuzzyResponse struct {
+	Success bool     `json:"success"`
+	Hashes  []string `json:"hashes"`
+}
+
 // PingResponse encapsulates the response of Ping.
 type PingResponse string
 
@@ -60,56 +100,72 @@ type UnexpectedResponseError struct {
 // New returns a client.
 // It takes the url of a rspamd instance, and configures the client with Options which are closures.
 func New(url string, options ...Option) *client {
-	client := &client{
-		client: resty.New().SetHostURL(url),
-	}
+	cl := NewFromClient(resty.New().SetBaseURL(url))
 
 	for _, option := range options {
-		err := option(client)
+		err := option(cl)
 		if err != nil {
 			log.Fatal("failed to configure client")
 		}
 	}
 
-	return client
+	return cl
+}
+
+// NewFromClient returns a client.
+// It takes an instance of resty.Client.
+func NewFromClient(restyClient *resty.Client) *client {
+	cl := &client{
+		client: restyClient,
+	}
+	return cl
 }
 
 // Check scans an email, returning a spam score and list of symbols.
-func (c *client) Check(ctx context.Context, e *Email) (*CheckResponse, error) {
+func (c *client) Check(ctx context.Context, cr *CheckRequest) (*CheckResponse, error) {
 	result := &CheckResponse{}
-	req := c.makeEmailRequest(ctx, e).SetResult(result)
+	req := c.buildRequest(ctx, cr.Message, cr.Header).SetResult(result)
 	_, err := c.sendRequest(req, resty.MethodPost, checkV2Endpoint)
 	return result, err
 }
 
 // LearnSpam trains rspamd's Bayesian classifier by marking an email as spam.
-func (c *client) LearnSpam(ctx context.Context, e *Email) (*LearnResponse, error) {
+func (c *client) LearnSpam(ctx context.Context, lr *LearnRequest) (*LearnResponse, error) {
 	result := &LearnResponse{}
-	req := c.makeEmailRequest(ctx, e).SetResult(result)
+	req := c.buildRequest(ctx, lr.Message, lr.Header).SetResult(result)
 	_, err := c.sendRequest(req, resty.MethodPost, learnSpamEndpoint)
 	return result, err
 }
 
-// LearnSpam trains rspamd's Bayesian classifier by marking an email as ham.
-func (c *client) LearnHam(ctx context.Context, e *Email) (*LearnResponse, error) {
+// LearnHam trains rspamd's Bayesian classifier by marking an email as ham.
+func (c *client) LearnHam(ctx context.Context, lr *LearnRequest) (*LearnResponse, error) {
 	result := &LearnResponse{}
-	req := c.makeEmailRequest(ctx, e).SetResult(result)
+	req := c.buildRequest(ctx, lr.Message, lr.Header).SetResult(result)
 	_, err := c.sendRequest(req, resty.MethodPost, learnHamEndpoint)
 	return result, err
 }
 
 // FuzzyAdd adds an email to fuzzy storage.
-func (c *client) FuzzyAdd(ctx context.Context, e *Email) (*LearnResponse, error) {
-	result := &LearnResponse{}
-	req := c.makeEmailRequest(ctx, e).SetResult(result)
+func (c *client) FuzzyAdd(ctx context.Context, fr *FuzzyRequest) (*FuzzyResponse, error) {
+	result := &FuzzyResponse{}
+	if fr.Header == nil {
+		fr.Header = http.Header{}
+	}
+	SetFlag(fr.Header, fr.Flag)
+	SetWeight(fr.Header, fr.Weight)
+	req := c.buildRequest(ctx, fr.Message, fr.Header).SetResult(result)
 	_, err := c.sendRequest(req, resty.MethodPost, fuzzyAddEndpoint)
 	return result, err
 }
 
-// FuzzyAdd removes an email from fuzzy storage.
-func (c *client) FuzzyDel(ctx context.Context, e *Email) (*LearnResponse, error) {
-	result := &LearnResponse{}
-	req := c.makeEmailRequest(ctx, e).SetResult(result)
+// FuzzyDel removes an email from fuzzy storage.
+func (c *client) FuzzyDel(ctx context.Context, fr *FuzzyRequest) (*FuzzyResponse, error) {
+	result := &FuzzyResponse{}
+	if fr.Header == nil {
+		fr.Header = http.Header{}
+	}
+	SetFlag(fr.Header, fr.Flag)
+	req := c.buildRequest(ctx, fr.Message, fr.Header).SetResult(result)
 	_, err := c.sendRequest(req, resty.MethodPost, fuzzyDelEndpoint)
 	return result, err
 }
@@ -121,25 +177,15 @@ func (c *client) Ping(ctx context.Context) (PingResponse, error) {
 	return result, err
 }
 
-func (c *client) makeEmailRequest(ctx context.Context, e *Email) *resty.Request {
-	headers := map[string]string{}
-	if e.queueID != "" {
-		headers["Queue-ID"] = e.queueID
-	}
-	if e.options.flag != 0 {
-		headers["Flag"] = fmt.Sprintf("%d", e.options.flag)
-	}
-	if e.options.weight != 0 {
-		headers["Weight"] = fmt.Sprintf("%d", e.options.weight)
-	}
+func (c *client) buildRequest(ctx context.Context, message io.Reader, Header http.Header) *resty.Request {
 	return c.client.R().
 		SetContext(ctx).
-		SetHeaders(headers).
-		SetBody(e.message)
+		SetHeaderMultiValues(Header).
+		SetBody(message)
 }
 
-func (c *client) sendRequest(req *resty.Request, method, url string) (*resty.Response, error) {
-	res, err := req.Execute(method, url)
+func (c *client) sendRequest(req *resty.Request, method, endpoint string) (*resty.Response, error) {
+	res, err := req.Execute(method, endpoint)
 
 	if err != nil {
 		return nil, fmt.Errorf("executing request: %q", err)
@@ -176,3 +222,30 @@ func IsAlreadyLearnedError(err error) bool {
 	var errResp *UnexpectedResponseError
 	return errors.As(err, &errResp) && errResp.Status == http.StatusAlreadyReported
 }
+
+func ReaderFromWriterTo(writerTo io.WriterTo) io.Reader {
+	r, w := io.Pipe()
+
+	go func() {
+		if _, err := writerTo.WriteTo(w); err != nil {
+			_ = w.CloseWithError(fmt.Errorf("writing to pipe: %q", err))
+			return
+		}
+
+		_ = w.Close() // Always succeeds
+	}()
+
+	return r
+}
+
+func SetQueueID(header http.Header, queueID string) {
+	header.Set(QueueID, queueID)
+}
+
+func SetFlag(header http.Header, flag int) {
+	header.Set(Flag, strconv.Itoa(flag))
+}
+
+func SetWeight(header http.Header, weight int) {
+	header.Set(Weight, strconv.Itoa(weight))
+}
diff --git a/client_mock.go b/client_mock.go
@@ -19,49 +19,49 @@ type mockClient struct {
 	mock.Mock
 }
 
-func (m *mockClient) Check(ctx context.Context, e *Email) (*CheckResponse, error) {
-	args := m.Called(ctx, e)
+func (m *mockClient) Check(ctx context.Context, cr *CheckRequest) (*CheckResponse, error) {
+	args := m.Called(ctx, cr)
 	if args.Error(1) != nil {
 		return nil, args.Error(1)
 	}
 
 	return args.Get(0).(*CheckResponse), nil
 }
 
-func (m *mockClient) LearnSpam(ctx context.Context, e *Email) (*LearnResponse, error) {
-	args := m.Called(ctx, e)
+func (m *mockClient) LearnSpam(ctx context.Context, lr *LearnRequest) (*LearnResponse, error) {
+	args := m.Called(ctx, lr)
 	if args.Error(1) != nil {
 		return nil, args.Error(1)
 	}
 
 	return args.Get(0).(*LearnResponse), nil
 }
 
-func (m *mockClient) LearnHam(ctx context.Context, e *Email) (*LearnResponse, error) {
-	args := m.Called(ctx, e)
+func (m *mockClient) LearnHam(ctx context.Context, lr *LearnRequest) (*LearnResponse, error) {
+	args := m.Called(ctx, lr)
 	if args.Error(1) != nil {
 		return nil, args.Error(1)
 	}
 
 	return args.Get(0).(*LearnResponse), nil
 }
 
-func (m *mockClient) FuzzyAdd(ctx context.Context, e *Email) (*LearnResponse, error) {
-	args := m.Called(ctx, e)
+func (m *mockClient) FuzzyAdd(ctx context.Context, fr *FuzzyRequest) (*FuzzyResponse, error) {
+	args := m.Called(ctx, fr)
 	if args.Error(1) != nil {
 		return nil, args.Error(1)
 	}
 
-	return args.Get(0).(*LearnResponse), nil
+	return args.Get(0).(*FuzzyResponse), nil
 }
 
-func (m *mockClient) FuzzyDel(ctx context.Context, e *Email) (*LearnResponse, error) {
-	args := m.Called(ctx, e)
+func (m *mockClient) FuzzyDel(ctx context.Context, fr *FuzzyRequest) (*FuzzyResponse, error) {
+	args := m.Called(ctx, fr)
 	if args.Error(1) != nil {
 		return nil, args.Error(1)
 	}
 
-	return args.Get(0).(*LearnResponse), nil
+	return args.Get(0).(*FuzzyResponse), nil
 }
 
 func (m *mockClient) Ping(ctx context.Context) (PingResponse, error) {