diff --git a/ext/ext.go b/ext/ext.go
index 69c7a52be49b0064ad5ab78dd9b95031cca525d6..6a2eb8ffba3bb37603eb40e890d2600101067942 100644
--- a/ext/ext.go
+++ b/ext/ext.go
@@ -1,6 +1,14 @@
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)
+ }
+}
diff --git a/ext/geoip/geoip.go b/ext/geoip/geoip.go
index 9eac8a24fd8fe3fedd27812155a9c854c637f30c..470a75a3b370eeaa91e14fa4e49d1dcddd8ff90b 100644
--- a/ext/geoip/geoip.go
+++ b/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
}
diff --git a/script/script.go b/script/script.go
index 55f0be391f97619c5b5e8341e19afa4b4788f7ba..c55752b4c11e970e6b98416e4c72e1836decc7c7 100644
--- a/script/script.go
+++ b/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 {
diff --git a/script/script_test.go b/script/script_test.go
index 2a0b71adb25085f0a9398c8fea6ffd169dd1c3b3..3fc38784ef3a391841403d14946fd45ce3b79ddd 100644
--- a/script/script_test.go
+++ b/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,
+ },
+ })
}
diff --git a/script/types.go b/script/types.go
index 4e9d7a01b8cc02ba20098b6d29fa5741b2ac642e..cd69b6231eff6aa307db8f73ef3f44d6fc63ca93 100644
--- a/script/types.go
+++ b/script/types.go
@@ -1,8 +1,12 @@
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
+}
diff --git a/server/server.go b/server/server.go
index 9edf106a08619e824192f7698afd185a914184e2..3d82f86f5d4779cbb85ce68f3c0af3aecd7f7c13 100644
--- a/server/server.go
+++ b/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)
}