<commit_before>package broker

import (
	"errors"
	"io"

	pb_api "github.com/TheThingsNetwork/ttn/api"
	pb "github.com/TheThingsNetwork/ttn/api/broker"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

type brokerRPC struct {
	broker Broker
}

func getCallerFromMetadata(ctx context.Context) (callerID string, err error) {
	md, ok := metadata.FromContext(ctx)
	// TODO: Check OK
	id, ok := md["id"]
	if !ok || len(id) < 1 {
		err = errors.New("ttn/broker: Caller did not provide \"id\" in context")
		return
	}
	callerID = id[0]
	if err != nil {
		return
	}
	token, ok := md["token"]
	if !ok || len(token) < 1 {
		err = errors.New("ttn/broker: Caller did not provide \"token\" in context")
		return
	}
	if token[0] != "token" {
		// TODO: Validate Token
		err = errors.New("ttn/broker: Caller not authorized")
		return
	}

	return
}

func (b *brokerRPC) Associate(stream pb.Broker_AssociateServer) error {
	routerID, err := getCallerFromMetadata(stream.Context())
	if err != nil {
		return err
	}
	downlinkChannel, err := b.broker.ActivateRouter(routerID)
	if err != nil {
		return err
	}
	defer b.broker.DeactivateRouter(routerID)
	go func() {
		for {
			if downlinkChannel == nil {
				return
			}
			select {
			case <-stream.Context().Done():
				return
			case downlink := <-downlinkChannel:
				if downlink != nil {
					if err := stream.Send(downlink); err != nil {
						// TODO: Check if the stream should be closed here
						return
					}
				}
			}
		}
	}()
	for {
		uplink, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		go b.broker.HandleUplink(uplink)
	}
}

func (b *brokerRPC) Subscribe(req *pb.SubscribeRequest, stream pb.Broker_SubscribeServer) error {
	handlerID, err := getCallerFromMetadata(stream.Context())
	if err != nil {
		return err
	}
	uplinkChannel, err := b.broker.ActivateHandler(handlerID)
	if err != nil {
		return err
	}
	defer b.broker.DeactivateHandler(handlerID)
	for {
		if uplinkChannel == nil {
			return nil
		}
		select {
		case <-stream.Context().Done():
			return stream.Context().Err()
		case uplink := <-uplinkChannel:
			if uplink != nil {
				if err := stream.Send(uplink); err != nil {
					return err
				}
			}
		}
	}
}

func (b *brokerRPC) Publish(stream pb.Broker_PublishServer) error {
	handlerID, err := getCallerFromMetadata(stream.Context())
	if err != nil {
		return err
	}
	for {
		downlink, err := stream.Recv()
		if err == io.EOF {
			return stream.SendAndClose(&pb_api.Ack{})
		}
		if err != nil {
			return err
		}
		// TODO: Validate that this handler can publish downlink for the application
		_ = handlerID
		go b.broker.HandleDownlink(downlink)
	}
}

func (b *brokerRPC) Activate(ctx context.Context, req *pb.DeviceActivationRequest) (res *pb.DeviceActivationResponse, err error) {
	_, err = getCallerFromMetadata(ctx)
	if err != nil {
		return nil, err
	}
	return b.broker.HandleActivation(req)
}

func (b *broker) RegisterRPC(s *grpc.Server) {
	server := &brokerRPC{b}
	pb.RegisterBrokerServer(s, server)
}
<commit_msg>Use grpc codes in broker errors<commit_after>package broker

import (
	"io"

	pb_api "github.com/TheThingsNetwork/ttn/api"
	pb "github.com/TheThingsNetwork/ttn/api/broker"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
)

type brokerRPC struct {
	broker Broker
}

func getCallerFromMetadata(ctx context.Context) (callerID string, err error) {
	md, ok := metadata.FromContext(ctx)
	// TODO: Check OK
	id, ok := md["id"]
	if !ok || len(id) < 1 {
		err = grpc.Errorf(codes.Unauthenticated, "ttn/broker: Caller did not provide \"id\" in context")
		return
	}
	callerID = id[0]
	if err != nil {
		return
	}
	token, ok := md["token"]
	if !ok || len(token) < 1 {
		err = grpc.Errorf(codes.Unauthenticated, "ttn/broker: Caller did not provide \"token\" in context")
		return
	}
	if token[0] != "token" {
		// TODO: Validate Token
		err = grpc.Errorf(codes.Unauthenticated, "ttn/broker: Caller not authorized")
		return
	}

	return
}

func (b *brokerRPC) Associate(stream pb.Broker_AssociateServer) error {
	routerID, err := getCallerFromMetadata(stream.Context())
	if err != nil {
		return err
	}
	downlinkChannel, err := b.broker.ActivateRouter(routerID)
	if err != nil {
		return err
	}
	defer b.broker.DeactivateRouter(routerID)
	go func() {
		for {
			if downlinkChannel == nil {
				return
			}
			select {
			case <-stream.Context().Done():
				return
			case downlink := <-downlinkChannel:
				if downlink != nil {
					if err := stream.Send(downlink); err != nil {
						// TODO: Check if the stream should be closed here
						return
					}
				}
			}
		}
	}()
	for {
		uplink, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		go b.broker.HandleUplink(uplink)
	}
}

func (b *brokerRPC) Subscribe(req *pb.SubscribeRequest, stream pb.Broker_SubscribeServer) error {
	handlerID, err := getCallerFromMetadata(stream.Context())
	if err != nil {
		return err
	}
	uplinkChannel, err := b.broker.ActivateHandler(handlerID)
	if err != nil {
		return err
	}
	defer b.broker.DeactivateHandler(handlerID)
	for {
		if uplinkChannel == nil {
			return nil
		}
		select {
		case <-stream.Context().Done():
			return stream.Context().Err()
		case uplink := <-uplinkChannel:
			if uplink != nil {
				if err := stream.Send(uplink); err != nil {
					return err
				}
			}
		}
	}
}

func (b *brokerRPC) Publish(stream pb.Broker_PublishServer) error {
	handlerID, err := getCallerFromMetadata(stream.Context())
	if err != nil {
		return err
	}
	for {
		downlink, err := stream.Recv()
		if err == io.EOF {
			return stream.SendAndClose(&pb_api.Ack{})
		}
		if err != nil {
			return err
		}
		// TODO: Validate that this handler can publish downlink for the application
		_ = handlerID
		go b.broker.HandleDownlink(downlink)
	}
}

func (b *brokerRPC) Activate(ctx context.Context, req *pb.DeviceActivationRequest) (res *pb.DeviceActivationResponse, err error) {
	_, err = getCallerFromMetadata(ctx)
	if err != nil {
		return nil, err
	}
	return b.broker.HandleActivation(req)
}

func (b *broker) RegisterRPC(s *grpc.Server) {
	server := &brokerRPC{b}
	pb.RegisterBrokerServer(s, server)
}

{'fraction_non_alphanumeric': 0.13484220592923174, 'fraction_numerical': 0.001275103602167676, 'mean_word_length': 3.3819832402234637, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4400731', 'n_tokens_mistral': 2536, 'n_tokens_neox': 2297, 'n_words': 562}