From 05204437d67e0acb015cdf1cdaba47f46047122f Mon Sep 17 00:00:00 2001
From: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>
Date: Thu, 12 Feb 2026 18:36:04 +0900
Subject: [PATCH] profile: spilt parser to
 github.com/AkihiroSuda/gomoddirectivecomments

https://github.com/AkihiroSuda/gomoddirectivecomments

Signed-off-by: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>
---
 go.mod                             |   1 +
 go.sum                             |   2 +
 pkg/profile/fromgomod/fromgomod.go | 122 +++--------------------------
 3 files changed, 14 insertions(+), 111 deletions(-)

diff --git a/go.mod b/go.mod
index 2770e3c..d72c789 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module github.com/AkihiroSuda/gomodjail
 go 1.24.0
 
 require (
+	github.com/AkihiroSuda/gomoddirectivecomments v0.1.0
 	github.com/AkihiroSuda/gosocialcheck v0.0.4
 	github.com/elastic/go-seccomp-bpf v1.6.0
 	github.com/spf13/cobra v1.10.2
diff --git a/go.sum b/go.sum
index 5893c23..9d2c782 100644
--- a/go.sum
+++ b/go.sum
@@ -1,3 +1,5 @@
+github.com/AkihiroSuda/gomoddirectivecomments v0.1.0 h1:5sKxYIkq9GGs0DTnuPNVm2Z/LmhKdTN+8QblThzTKqg=
+github.com/AkihiroSuda/gomoddirectivecomments v0.1.0/go.mod h1:flXOhVLWfsi4FuFhFoc9F3m7wAH2RT4aM3fVyQROKt4=
 github.com/AkihiroSuda/gosocialcheck v0.0.4 h1:DSeLoaG7jgjbscgFfyRzdLQprpTnomVxJF7pRAdjRnw=
 github.com/AkihiroSuda/gosocialcheck v0.0.4/go.mod h1:DNDBTlEkGah2wiGC50XNFPETBNiN1o2c15w9ziECOtM=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
diff --git a/pkg/profile/fromgomod/fromgomod.go b/pkg/profile/fromgomod/fromgomod.go
index 4c5de65..5f1ccef 100644
--- a/pkg/profile/fromgomod/fromgomod.go
+++ b/pkg/profile/fromgomod/fromgomod.go
@@ -4,8 +4,8 @@ import (
 	"fmt"
 	"log/slog"
 	"slices"
-	"strings"
 
+	"github.com/AkihiroSuda/gomoddirectivecomments"
 	"golang.org/x/mod/modfile"
 
 	"github.com/AkihiroSuda/gomodjail/pkg/profile"
@@ -13,119 +13,19 @@ import (
 
 func FromGoMod(mod *modfile.File, prof *profile.Profile) error {
 	prof.Module = mod.Module.Mod.Path
-	currentDefaultPolicy := profile.PolicyUnconfined
-
-	for _, c := range append(mod.Module.Syntax.Before, mod.Module.Syntax.Suffix...) {
-		if tok := c.Token; tok != "" {
-			pol, err := policyFromComment(tok)
-			if err != nil {
-				err = fmt.Errorf("failed to parse comment %+v: %w", c, err)
-				return err
-			}
-			currentDefaultPolicy = pol
-		}
+	var err error
+	parsedPolicies, err := gomoddirectivecomments.Parse(mod, "gomodjail", profile.PolicyUnconfined)
+	if err != nil {
+		return fmt.Errorf("failed to parse Go module directive comments: %w", err)
 	}
-
-	for _, c := range append(mod.Go.Syntax.Before, mod.Go.Syntax.Suffix...) {
-		if tok := c.Token; tok != "" {
-			pol, err := policyFromComment(tok)
-			if err != nil {
-				err = fmt.Errorf("failed to parse comment %+v: %w", c, err)
-				return err
-			}
-			return fmt.Errorf("policy %q is specified in an invalid position", pol)
+	for modPath, pol := range parsedPolicies {
+		if !slices.Contains(profile.KnownPolicies, pol) {
+			return fmt.Errorf("module %q: unknown policy %q", modPath, pol)
 		}
-	}
-
-	for _, f := range mod.Require {
-		if syn := f.Syntax; syn != nil {
-			pol := currentDefaultPolicy
-			if syn.InBlock {
-				// TODO: cache line blocks
-				if lineBlock := findLineBlock(mod.Syntax.Stmt, syn); lineBlock != nil {
-					lineBlockPol, err := policyFromLineBlock(lineBlock)
-					if err != nil {
-						err = fmt.Errorf("failed to parse line block %+v: %w", lineBlock, err)
-						return err
-					}
-					if lineBlockPol != "" {
-						pol = lineBlockPol
-					}
-				}
-			}
-			for _, c := range append(syn.Before, syn.Suffix...) {
-				if tok := c.Token; tok != "" {
-					polFromComment, err := policyFromComment(tok)
-					if err != nil {
-						err = fmt.Errorf("failed to parse comment %+v: %w", c, err)
-						return err
-					}
-					if polFromComment != "" {
-						pol = polFromComment
-					}
-				}
-			}
-			if pol == "" {
-				pol = currentDefaultPolicy
-			}
-			if pol == profile.PolicyUnconfined {
-				pol = "" // reduce map size
-			}
-			if existPol, ok := prof.Modules[f.Mod.Path]; ok && existPol != pol {
-				slog.Warn("Overwriting an existing policy", "module", f.Mod.Path, "old", existPol, "new", pol)
-			}
-			if pol == "" {
-				delete(prof.Modules, f.Mod.Path)
-			} else {
-				prof.Modules[f.Mod.Path] = pol
-			}
+		if existPol, ok := prof.Modules[modPath]; ok && existPol != pol {
+			slog.Warn("Overwriting an existing policy", "module", modPath, "old", existPol, "new", pol)
 		}
 	}
+	prof.Modules = parsedPolicies
 	return nil
 }
-
-func policyFromComment(token string) (string, error) {
-	token = strings.TrimPrefix(token, "//")
-	// TODO: support /* ... */
-	for _, f := range strings.Fields(token) {
-		f = strings.TrimPrefix(f, "//")
-		if strings.HasPrefix(f, "gomodjail:") {
-			pol := profile.Policy(strings.TrimPrefix(f, "gomodjail:"))
-			if !slices.Contains(profile.KnownPolicies, pol) {
-				return pol, fmt.Errorf("unknown policy %q", pol)
-			}
-			return pol, nil
-		}
-	}
-	return "", nil
-}
-
-func findLineBlock(exprs []modfile.Expr, line modfile.Expr) *modfile.LineBlock {
-	start, end := line.Span()
-	for _, expr := range exprs {
-		lb, ok := expr.(*modfile.LineBlock)
-		if !ok {
-			continue
-		}
-		lbStart, lbEnd := lb.Span()
-		if start.Line >= lbStart.Line && end.Line <= lbEnd.Line {
-			return lb
-		}
-	}
-	return nil
-}
-
-func policyFromLineBlock(lb *modfile.LineBlock) (string, error) {
-	for _, c := range append(lb.Before, lb.Suffix...) {
-		if tok := c.Token; tok != "" {
-			pol, err := policyFromComment(tok)
-			if err != nil {
-				return "", err
-			}
-			if pol != "" {
-				return pol, nil
-			}
-		}
-	}
-	return "", nil
-}