Add external source support to the Server type

ext/ext.go

 package ext
 
-import "github.com/d5/tengo/objects"
+import (
+	"errors"
+	"fmt"
+
+	"github.com/d5/tengo/objects"
+
+	"git.autistici.org/ai3/tools/iprep/ext/dnsbl"
+	"git.autistici.org/ai3/tools/iprep/ext/geoip"
+)
 
 // An ExternalSource provides per-IP information from third-party
 // sources. The lookup can return any Tengo object, we don't want to
@@ -9,3 +17,32 @@ import "github.com/d5/tengo/objects"
 type ExternalSource interface {
 	LookupIP(string) (objects.Object, error)
 }
+
+// New creates a new ExternalSource.
+func New(sourceType string, params map[string]interface{}) (ExternalSource, error) {
+	switch sourceType {
+
+	case "dnsbl":
+		domain, ok := params["domain"].(string)
+		if !ok {
+			return nil, errors.New("missing parameter 'domain'")
+		}
+		match, ok := params["match"].(string)
+		if !ok {
+			return nil, errors.New("missing parameter 'match'")
+		}
+		return dnsbl.New(domain, match)
+
+	case "geoip":
+		var paths []string
+		if l, ok := params["paths"].([]interface{}); ok {
+			for _, p := range l {
+				paths = append(paths, p.(string))
+			}
+		}
+		return geoip.New(paths)
+
+	default:
+		return nil, fmt.Errorf("unknown source type '%s'", sourceType)
+	}
+}

ext/geoip/geoip.go

@@ -15,7 +15,7 @@ type GeoIP struct {
 	readers []*maxminddb.Reader
 }
 
-func NewGeoIP(paths []string) (*GeoIP, error) {
+func New(paths []string) (*GeoIP, error) {
 	if len(paths) == 0 {
 		paths = defaultGeoIPPaths
 	}

script/script.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/d5/tengo/script"
 	"github.com/d5/tengo/stdlib"
+
+	"git.autistici.org/ai3/tools/iprep/ext"
 )
 
 // A Script is a compiled Tengo script that will be executed on every
@@ -72,7 +74,7 @@ func NewScript(src []byte) (*Script, error) {
 
 // RunIP evaluates the script once with the provided context and
 // returns the resulting score.
-func (script *Script) RunIP(ctx context.Context, ip string, counts map[string]int64, intervalSecs float64, ext map[string]interface{}) (float64, error) {
+func (script *Script) RunIP(ctx context.Context, ip string, counts map[string]int64, intervalSecs float64, extSrcs map[string]ext.ExternalSource) (float64, error) {
 	c := script.compiled.Clone()
 
 	// Set the global variables that constitute the script
@@ -88,7 +90,7 @@ func (script *Script) RunIP(ctx context.Context, ip string, counts map[string]in
 	if err := c.Set("interval", intervalSecs); err != nil {
 		return 0, err
 	}
-	if err := c.Set("ext", ext); err != nil {
+	if err := c.Set("ext", extLookupFunc(extSrcs)); err != nil {
 		return 0, err
 	}
 	if err := c.Set("counts", intMap(counts)); err != nil {

script/script_test.go

@@ -3,9 +3,23 @@ package script
 import (
 	"context"
 	"testing"
+
+	"github.com/d5/tengo/objects"
+
+	"git.autistici.org/ai3/tools/iprep/ext"
 )
 
-func runTestScript(t *testing.T, src string, expected float64) {
+type dummyExtSource map[string]int64
+
+func (s dummyExtSource) LookupIP(ip string) (objects.Object, error) {
+	value, ok := s[ip]
+	if !ok {
+		return objects.UndefinedValue, nil
+	}
+	return &objects.Int{Value: value}, nil
+}
+
+func runTestScript(t *testing.T, src string, expected float64, extSrcs map[string]ext.ExternalSource) {
 	s, err := NewScript([]byte(src))
 	if err != nil {
 		t.Fatalf("NewScript: %v", err)
@@ -15,7 +29,7 @@ func runTestScript(t *testing.T, src string, expected float64) {
 		"test2": 2,
 	}
 
-	score, err := s.RunIP(context.Background(), "1.2.3.4", m, 3600, nil)
+	score, err := s.RunIP(context.Background(), "1.2.3.4", m, 3600, extSrcs)
 	if err != nil {
 		t.Fatalf("runScript: %v", err)
 	}
@@ -28,23 +42,33 @@ func TestScript_WithLookup(t *testing.T) {
 	runTestScript(t, `
 score = counts["test"] / 2 + counts["test2"]
 //score = 7
-`, 7)
+`, 7, nil)
 }
 
 func TestScript_FloatScore(t *testing.T) {
 	runTestScript(t, `
 score = 7.0
-`, 7)
+`, 7, nil)
 }
 
 func TestScript_IntScore(t *testing.T) {
 	runTestScript(t, `
 score = 7
-`, 7)
+`, 7, nil)
 }
 
 func TestScript_StringScore(t *testing.T) {
 	runTestScript(t, `
 score = "7"
-`, 7)
+`, 7, nil)
+}
+
+func TestScript_ExternalSource(t *testing.T) {
+	runTestScript(t, `
+score = ext("test", ip)
+`, 7, map[string]ext.ExternalSource{
+		"test": dummyExtSource{
+			"1.2.3.4": 7,
+		},
+	})
 }

script/types.go

 package script
 
 import (
+	"log"
+
 	"github.com/d5/tengo/compiler/token"
 	"github.com/d5/tengo/objects"
+
+	"git.autistici.org/ai3/tools/iprep/ext"
 )
 
 // Tengo object type that wraps a read-only map[string]int64, without
@@ -105,3 +109,70 @@ func (i *intMapIterator) Key() objects.Object {
 func (i *intMapIterator) Value() objects.Object {
 	return &objects.Int{Value: i.arr[i.idx-1].value}
 }
+
+// A Tengo callable to look up data in external sources.
+type extLookupFunc map[string]ext.ExternalSource
+
+func (f extLookupFunc) String() string {
+	return "<extLookupFunc>"
+}
+
+func (f extLookupFunc) TypeName() string {
+	return "ext-lookup-func"
+}
+
+func (f extLookupFunc) Copy() objects.Object {
+	return f
+}
+
+func (f extLookupFunc) IsFalsy() bool {
+	return len(f) == 0
+}
+
+func (f extLookupFunc) Equals(o objects.Object) bool {
+	return false
+}
+
+func (f extLookupFunc) BinaryOp(op token.Token, rhs objects.Object) (objects.Object, error) {
+	return nil, objects.ErrInvalidOperator
+}
+
+func (f extLookupFunc) Call(args ...objects.Object) (ret objects.Object, err error) {
+	// Expected arguments: source name, IP.
+	if len(args) != 2 {
+		return nil, objects.ErrWrongNumArguments
+	}
+
+	name, ok := objects.ToString(args[0])
+	if !ok {
+		return nil, objects.ErrInvalidArgumentType{
+			Name:     "source_name",
+			Expected: "string",
+			Found:    args[0].TypeName(),
+		}
+	}
+
+	ip, ok := objects.ToString(args[1])
+	if !ok {
+		return nil, objects.ErrInvalidArgumentType{
+			Name:     "ip",
+			Expected: "string",
+			Found:    args[1].TypeName(),
+		}
+	}
+
+	// Invoke the external source lookup method.
+	if f == nil {
+		return nil, objects.ErrInvalidIndexValueType
+	}
+	src, ok := f[name]
+	if !ok {
+		return nil, objects.ErrInvalidIndexValueType
+	}
+	result, err := src.LookupIP(ip)
+	if err != nil {
+		log.Printf("external source lookup error: ip=%s: %v", ip, err)
+		return objects.UndefinedValue, nil
+	}
+	return result, nil
+}

server/server.go

@@ -10,12 +10,14 @@ import (
 	"google.golang.org/grpc/status"
 
 	"git.autistici.org/ai3/tools/iprep/db"
+	"git.autistici.org/ai3/tools/iprep/ext"
 	ippb "git.autistici.org/ai3/tools/iprep/proto"
 	"git.autistici.org/ai3/tools/iprep/script"
 )
 
 type Server struct {
 	*script.Manager
+	ext     map[string]ext.ExternalSource
 	horizon time.Duration
 	db      db.DB
 	stop    chan struct{}
@@ -30,7 +32,7 @@ var (
 	defaultScript = `score = 0`
 )
 
-func New(dbPath, scriptPath string) (*Server, error) {
+func New(dbPath, scriptPath string, extSrcs map[string]ext.ExternalSource) (*Server, error) {
 	scriptMgr, err := script.NewManager(scriptPath, defaultScript)
 	if err != nil {
 		return nil, err
@@ -43,6 +45,7 @@ func New(dbPath, scriptPath string) (*Server, error) {
 
 	s := &Server{
 		Manager: scriptMgr,
+		ext:     extSrcs,
 		horizon: defaultHorizon,
 		db:      database,
 		stop:    make(chan struct{}),
@@ -88,7 +91,7 @@ func (s *Server) GetScore(ctx context.Context, req *ippb.GetScoreRequest) (*ippb
 		return nil, status.Errorf(codes.Unavailable, "%v", err)
 	}
 
-	score, err := s.Script().RunIP(ctx, req.Ip, counts, s.horizon.Seconds(), nil)
+	score, err := s.Script().RunIP(ctx, req.Ip, counts, s.horizon.Seconds(), s.ext)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "%v", err)
 	}