grpc: support expire/delete API keys by ID

Update ExpireApiKey and DeleteApiKey handlers to accept either ID or
prefix for identifying the API key. Returns InvalidArgument error if
neither or both are provided.

Add tests for:
- Expire by ID
- Expire by prefix (backwards compatibility)
- Delete by ID
- Delete by prefix (backwards compatibility)
- Error when neither ID nor prefix provided
- Error when both ID and prefix provided

Updates #2986

hscontrol/grpcv1.go

@@ -577,14 +577,35 @@ func (api headscaleV1APIServer) CreateApiKey(
 	return &v1.CreateApiKeyResponse{ApiKey: apiKey}, nil
 }
 
+// apiKeyIdentifier is implemented by requests that identify an API key.
+type apiKeyIdentifier interface {
+	GetId() uint64
+	GetPrefix() string
+}
+
+// getAPIKey retrieves an API key by ID or prefix from the request.
+// Returns InvalidArgument if neither or both are provided.
+func (api headscaleV1APIServer) getAPIKey(req apiKeyIdentifier) (*types.APIKey, error) {
+	hasID := req.GetId() != 0
+	hasPrefix := req.GetPrefix() != ""
+
+	switch {
+	case hasID && hasPrefix:
+		return nil, status.Error(codes.InvalidArgument, "provide either id or prefix, not both")
+	case hasID:
+		return api.h.state.GetAPIKeyByID(req.GetId())
+	case hasPrefix:
+		return api.h.state.GetAPIKey(req.GetPrefix())
+	default:
+		return nil, status.Error(codes.InvalidArgument, "must provide id or prefix")
+	}
+}
+
 func (api headscaleV1APIServer) ExpireApiKey(
 	ctx context.Context,
 	request *v1.ExpireApiKeyRequest,
 ) (*v1.ExpireApiKeyResponse, error) {
-	var apiKey *types.APIKey
-	var err error
-
-	apiKey, err = api.h.state.GetAPIKey(request.Prefix)
+	apiKey, err := api.getAPIKey(request)
 	if err != nil {
 		return nil, err
 	}
@@ -622,12 +643,7 @@ func (api headscaleV1APIServer) DeleteApiKey(
 	ctx context.Context,
 	request *v1.DeleteApiKeyRequest,
 ) (*v1.DeleteApiKeyResponse, error) {
-	var (
-		apiKey *types.APIKey
-		err    error
-	)
-
-	apiKey, err = api.h.state.GetAPIKey(request.Prefix)
+	apiKey, err := api.getAPIKey(request)
 	if err != nil {
 		return nil, err
 	}