2 files changed, 46 insertions, 36 deletions

diff --git a/server/grpc_server.go b/server/grpc_server.go
index d68ecd96..44c3562d 100644
--- a/server/grpc_server.go
+++ b/server/grpc_server.go
@@ -63,6 +63,7 @@ const (
 	REQ_MOD_DEFINED_SET
 	REQ_STATEMENT
 	REQ_MOD_STATEMENT
+	REQ_POLICY
 )
 
 const GRPC_PORT = 8080
@@ -326,42 +327,6 @@ func (s *Server) modPolicy(arg *api.PolicyArguments, stream interface{}) error {
 	return nil
 }
 
-func (s *Server) GetPolicyRoutePolicies(arg *api.PolicyArguments, stream api.GobgpApi_GetPolicyRoutePoliciesServer) error {
-	var rf bgp.RouteFamily
-	var reqType int
-	switch arg.Resource {
-	case api.Resource_POLICY_ROUTEPOLICY:
-		reqType = REQ_POLICY_ROUTEPOLICIES
-	default:
-		return fmt.Errorf("unsupported resource type: %v", arg.Resource)
-	}
-	req := NewGrpcRequest(reqType, "", rf, arg)
-	s.bgpServerCh <- req
-	return handleMultipleResponses(req, func(res *GrpcResponse) error {
-		return stream.Send(res.Data.(*api.PolicyDefinition))
-	})
-}
-
-func (s *Server) GetPolicyRoutePolicy(ctx context.Context, arg *api.PolicyArguments) (*api.PolicyDefinition, error) {
-	var rf bgp.RouteFamily
-	var reqType int
-	switch arg.Resource {
-	case api.Resource_POLICY_ROUTEPOLICY:
-		reqType = REQ_POLICY_ROUTEPOLICY
-	default:
-		return nil, fmt.Errorf("unsupported resource type: %v", arg.Resource)
-	}
-	req := NewGrpcRequest(reqType, "", rf, arg)
-	s.bgpServerCh <- req
-
-	res := <-req.ResponseCh
-	if err := res.Err(); err != nil {
-		log.Debug(err.Error())
-		return nil, err
-	}
-	return res.Data.(*api.PolicyDefinition), nil
-}
-
 func (s *Server) ModPolicyRoutePolicy(stream api.GobgpApi_ModPolicyRoutePolicyServer) error {
 	for {
 		arg, err := stream.Recv()
@@ -486,6 +451,22 @@ func (s *Server) ModStatement(ctx context.Context, arg *api.ModStatementArgument
 	return s.mod(REQ_MOD_STATEMENT, arg)
 }
 
+func (s *Server) GetPolicy(ctx context.Context, arg *api.Policy) (*api.Policy, error) {
+	d, err := s.get(REQ_POLICY, arg)
+	if err != nil {
+		return nil, err
+	}
+	return d.(*api.Policy), nil
+}
+
+func (s *Server) GetPolicies(arg *api.Policy, stream api.GobgpApi_GetPoliciesServer) error {
+	req := NewGrpcRequest(REQ_POLICY, "", bgp.RouteFamily(0), arg)
+	s.bgpServerCh <- req
+	return handleMultipleResponses(req, func(res *GrpcResponse) error {
+		return stream.Send(res.Data.(*api.Policy))
+	})
+}
+
 type GrpcRequest struct {
 	RequestType int
 	Name        string
diff --git a/server/server.go b/server/server.go
index 69c27f14..91466753 100644
--- a/server/server.go
+++ b/server/server.go
@@ -1656,6 +1656,13 @@ func (server *BgpServer) handleGrpc(grpcReq *GrpcRequest) []*SenderMsg {
 			ResponseErr: err,
 		}
 		close(grpcReq.ResponseCh)
+	case REQ_POLICY:
+		if err := server.handleGrpcGetPolicy(grpcReq); err != nil {
+			grpcReq.ResponseCh <- &GrpcResponse{
+				ResponseErr: err,
+			}
+		}
+		close(grpcReq.ResponseCh)
 	case REQ_POLICY_ROUTEPOLICY, REQ_POLICY_ROUTEPOLICIES:
 		info := server.policy.PolicyMap
 		typ := grpcReq.RequestType
@@ -1819,6 +1826,28 @@ func (server *BgpServer) handleGrpcModStatement(grpcReq *GrpcRequest) error {
 
 }
 
+func (server *BgpServer) handleGrpcGetPolicy(grpcReq *GrpcRequest) error {
+	arg := grpcReq.Data.(*api.Policy)
+	name := arg.Name
+	found := false
+	for _, s := range server.policy.PolicyMap {
+		if name != "" && name != s.Name() {
+			continue
+		}
+		grpcReq.ResponseCh <- &GrpcResponse{
+			Data: s.ToApiStruct(),
+		}
+		found = true
+		if name != "" {
+			break
+		}
+	}
+	if !found {
+		return fmt.Errorf("not found %s", name)
+	}
+	return nil
+}
+
 func (server *BgpServer) handleMrt(grpcReq *GrpcRequest) {
 	now := uint32(time.Now().Unix())
 	view := ""