bgpwtf/cccampix/irr: limit concurrency

Change-Id: I958322f33c86469f9c3e21d1bd962faede2a3fee

bgpwtf/cccampix/irr/main.go

@@ -29,10 +29,10 @@ func main() {
 	}
 
 	s := &service{
-		iana: provider.NewIANA(),
+		iana: provider.NewIANA(2),
 		providers: map[provider.IRR]provider.Provider{
-			provider.IRR_RIPE: provider.NewRIPE(),
-			provider.IRR_ARIN: provider.NewARIN(),
+			provider.IRR_RIPE: provider.NewRIPE(10),
+			provider.IRR_ARIN: provider.NewARIN(2),
 		},
 	}
 	pb.RegisterIRRServer(mi.GRPC(), s)
@@ -68,7 +68,7 @@ func (s *service) Query(ctx context.Context, req *pb.IRRQueryRequest) (*pb.IRRQu
 
 	prov, ok := s.providers[irr]
 	if !ok {
-		return nil, status.Errorf(codes.NotFound, "AS belongs to unhandled IRR %s", irr.String())
+		return nil, status.Errorf(codes.Unimplemented, "AS belongs to unhandled IRR %s", irr.String())
 	}
 
 	res, err := prov.Query(ctx, asn)

bgpwtf/cccampix/irr/provider/BUILD.bazel

@@ -16,6 +16,8 @@ go_library(
         "//bgpwtf/cccampix/proto:go_default_library",
         "@com_github_golang_collections_go_datastructures//augmentedtree:go_default_library",
         "@com_github_golang_glog//:go_default_library",
+        "@org_golang_google_grpc//codes:go_default_library",
+        "@org_golang_google_grpc//status:go_default_library",
     ],
 )
 

bgpwtf/cccampix/irr/provider/arin.go

@@ -12,21 +12,31 @@ import (
 
 	"code.hackerspace.pl/hscloud/bgpwtf/cccampix/irr/whois"
 	pb "code.hackerspace.pl/hscloud/bgpwtf/cccampix/proto"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 const ARINWhois = "rr.arin.net:43"
 
 type arin struct {
+	sem chan struct{}
 }
 
-func NewARIN() Provider {
-	return &arin{}
+func NewARIN(limit int) Provider {
+	return &arin{
+		sem: make(chan struct{}, limit),
+	}
 }
 
-func (r *arin) Query(ctx context.Context, asn uint64) (*pb.IRRQueryResponse, error) {
+func (a *arin) Query(ctx context.Context, asn uint64) (*pb.IRRQueryResponse, error) {
+	a.sem <- struct{}{}
+	defer func() {
+		<-a.sem
+	}()
+
 	data, err := whois.Query(ctx, ARINWhois, fmt.Sprintf("AS%d", asn))
 	if err != nil {
-		return nil, fmt.Errorf("could not contact ARIN IRR: %v", err)
+		return nil, status.Errorf(codes.Unavailable, "could not contact ARIN IRR: %v", err)
 	}
 
 	lines := strings.Split(data, "\n")
@@ -53,14 +63,14 @@ func (r *arin) Query(ctx context.Context, asn uint64) (*pb.IRRQueryResponse, err
 		if strings.HasPrefix(line, " ") {
 			// Continuation
 			if len(attrs) < 1 {
-				return nil, fmt.Errorf("unparseable IRR, continuation with no previous atribute name: %q", line)
+				return nil, status.Errorf(codes.Unavailable, "unparseable IRR, continuation with no previous atribute name: %q", line)
 			}
 
 			attrs[len(attrs)-1].value += " " + strings.TrimSpace(line)
 		} else {
 			parts := strings.SplitN(line, ":", 2)
 			if len(parts) != 2 {
-				return nil, fmt.Errorf("unparseable IRR, line with no attribute key: %q", line)
+				return nil, status.Errorf(codes.Unavailable, "unparseable IRR, line with no attribute key: %q", line)
 			}
 			name := strings.TrimSpace(parts[0])
 			value := strings.TrimSpace(parts[1])
@@ -71,6 +81,10 @@ func (r *arin) Query(ctx context.Context, asn uint64) (*pb.IRRQueryResponse, err
 		}
 	}
 
+	if len(attrs) == 0 {
+		return nil, status.Errorf(codes.NotFound, "no such ASN")
+	}
+
 	return &pb.IRRQueryResponse{
 		Source:     pb.IRRQueryResponse_SOURCE_ARIN,
 		Attributes: parseAttributes(attrs),

bgpwtf/cccampix/irr/provider/iana.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/golang-collections/go-datastructures/augmentedtree"
 	"github.com/golang/glog"
@@ -39,9 +40,11 @@ type IANA struct {
 	// The tree library needs intervals to have a unique ID. We use a counter
 	// for this effect.
 	id uint64
+
+	mu sync.Mutex
 }
 
-func NewIANA() *IANA {
+func NewIANA(limit int) *IANA {
 	return &IANA{
 		cache: augmentedtree.New(1),
 	}
@@ -93,6 +96,9 @@ func (d *delegation) ID() uint64 {
 
 // Who returns the responsible IRR (or UNKNOWN) for a given AS.
 func (i *IANA) Who(ctx context.Context, asn uint64) (IRR, error) {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
 	q := &delegation{
 		id:   i.nextID(),
 		low:  int64(asn),
@@ -105,6 +111,7 @@ func (i *IANA) Who(ctx context.Context, asn uint64) (IRR, error) {
 	}
 
 	// No cache entry, query whois.
+
 	glog.Infof("Cache miss for AS%d", asn)
 	data, err := whois.Query(ctx, "whois.iana.org:43", fmt.Sprintf("AS%d", asn))
 	if err != nil {

bgpwtf/cccampix/irr/provider/ripe.go

@@ -11,6 +11,8 @@ import (
 	"net/http"
 
 	pb "code.hackerspace.pl/hscloud/bgpwtf/cccampix/proto"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 type ripeResponse struct {
@@ -32,13 +34,21 @@ type ripeAttribute struct {
 }
 
 type ripe struct {
+	sem chan struct{}
 }
 
-func NewRIPE() Provider {
-	return &ripe{}
+func NewRIPE(limit int) Provider {
+	return &ripe{
+		sem: make(chan struct{}, limit),
+	}
 }
 
 func (r *ripe) Query(ctx context.Context, as uint64) (*pb.IRRQueryResponse, error) {
+	r.sem <- struct{}{}
+	defer func() {
+		<-r.sem
+	}()
+
 	req, err := http.NewRequest("GET", fmt.Sprintf("http://rest.db.ripe.net/ripe/aut-num/AS%d.json", as), nil)
 	if err != nil {
 		return nil, err
@@ -49,22 +59,22 @@ func (r *ripe) Query(ctx context.Context, as uint64) (*pb.IRRQueryResponse, erro
 
 	res, err := client.Do(req)
 	if err != nil {
-		return nil, fmt.Errorf("could not run GET to RIPE: %v", err)
+		return nil, status.Errorf(codes.Unavailable, "could not run GET to RIPE: %v", err)
 	}
 	defer res.Body.Close()
 	bytes, err := ioutil.ReadAll(res.Body)
 	if err != nil {
-		return nil, fmt.Errorf("could not read response from RIPE: %v", err)
+		return nil, status.Errorf(codes.Unavailable, "could not read response from RIPE: %v", err)
 	}
 
 	data := ripeResponse{}
 	err = json.Unmarshal(bytes, &data)
 	if err != nil {
-		return nil, fmt.Errorf("could not decode response from RIPE: %v", err)
+		return nil, status.Errorf(codes.Unavailable, "could not decode response from RIPE: %v", err)
 	}
 
 	if len(data.Objects.Object) != 1 {
-		return nil, fmt.Errorf("could not retriev aut-num from RIPE")
+		return nil, status.Error(codes.NotFound, "could not retrieve aut-num from RIPE")
 	}
 
 	attributes := make([]rpslRawAttribute, len(data.Objects.Object[0].Attributes.Attribute))