mcp: token: handle authorization_code (pt2) (#5589)

internal/mcp/storage.go

@@ -7,10 +7,12 @@ import (
 	"github.com/google/uuid"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/timestamppb"
 
 	oauth21proto "github.com/pomerium/pomerium/internal/oauth21/gen"
 	rfc7591v1 "github.com/pomerium/pomerium/internal/rfc7591"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/grpc/session"
 	"github.com/pomerium/pomerium/pkg/protoutil"
 )
 
@@ -106,3 +108,40 @@ func (storage *Storage) GetAuthorizationRequest(
 
 	return v, nil
 }
+
+func (storage *Storage) DeleteAuthorizationRequest(
+	ctx context.Context,
+	id string,
+) error {
+	data := protoutil.NewAny(&oauth21proto.AuthorizationRequest{})
+	_, err := storage.client.Put(ctx, &databroker.PutRequest{
+		Records: []*databroker.Record{{
+			Id:        id,
+			Data:      data,
+			Type:      data.TypeUrl,
+			DeletedAt: timestamppb.Now(),
+		}},
+	})
+	if err != nil {
+		return fmt.Errorf("failed to delete authorization request by ID: %w", err)
+	}
+	return nil
+}
+
+func (storage *Storage) GetSession(ctx context.Context, id string) (*session.Session, error) {
+	v := new(session.Session)
+	rec, err := storage.client.Get(ctx, &databroker.GetRequest{
+		Type: protoutil.GetTypeURL(v),
+		Id:   id,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get session by ID: %w", err)
+	}
+
+	err = anypb.UnmarshalTo(rec.Record.Data, v, proto.UnmarshalOptions{})
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal session: %w", err)
+	}
+
+	return v, nil
+}