1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  //
  3  // Licensed under the Apache License, Version 2.0 (the "License");
  4  // you may not use this file except in compliance with the License.
  5  // You may obtain a copy of the License at
  6  //
  7  //     http://www.apache.org/licenses/LICENSE-2.0
  8  //
  9  // Unless required by applicable law or agreed to in writing, software
 10  // distributed under the License is distributed on an "AS IS" BASIS,
 11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  // See the License for the specific language governing permissions and
 13  // limitations under the License.
 14  //
 15  // Requirement: Any integration or derivative work must explicitly attribute
 16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  // documentation or user interface, as detailed in the NOTICE file.
 18  
 19  package websocket
 20  
 21  import (
 22  	"os"
 23  	"path/filepath"
 24  	"strings"
 25  	"testing"
 26  )
 27  
 28  // buildFakeClone creates a fake cloned repository directory tree in tmp:
 29  //
 30  //	<tmp>/data/fingerprints/foo.yaml
 31  //	<tmp>/data/vuln/bar/CVE-2024-0001.yaml
 32  //	<tmp>/data/mcp/tool.yaml
 33  //	<tmp>/README.md   <- should NOT be copied
 34  func buildFakeClone(t *testing.T, root string) {
 35  	t.Helper()
 36  	files := map[string]string{
 37  		"data/fingerprints/foo.yaml":       "name: foo\n",
 38  		"data/vuln/bar/CVE-2024-0001.yaml": "cve: CVE-2024-0001\n",
 39  		"data/mcp/tool.yaml":               "rule: test\n",
 40  		"README.md":                        "# readme\n",
 41  	}
 42  	for rel, content := range files {
 43  		path := filepath.Join(root, rel)
 44  		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
 45  			t.Fatalf("MkdirAll %s: %v", filepath.Dir(path), err)
 46  		}
 47  		if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
 48  			t.Fatalf("WriteFile %s: %v", path, err)
 49  		}
 50  	}
 51  }
 52  
 53  func TestCopyDataDirs_selectiveDirs(t *testing.T) {
 54  	srcRoot := t.TempDir()
 55  	buildFakeClone(t, srcRoot)
 56  
 57  	dstRoot := t.TempDir()
 58  	orig, _ := os.Getwd()
 59  	if err := os.Chdir(dstRoot); err != nil {
 60  		t.Fatalf("Chdir: %v", err)
 61  	}
 62  	defer os.Chdir(orig)
 63  
 64  	dirs := []string{"fingerprints", "vuln"}
 65  	n, err := copyDataDirs(srcRoot, dirs)
 66  	if err != nil {
 67  		t.Fatalf("copyDataDirs: %v", err)
 68  	}
 69  
 70  	// Expect 2 files: foo.yaml and CVE-2024-0001.yaml
 71  	if n != 2 {
 72  		t.Errorf("expected 2 files written, got %d", n)
 73  	}
 74  
 75  	// Verify fingerprints file exists with correct content.
 76  	fpPath := filepath.Join("data", "fingerprints", "foo.yaml")
 77  	data, err := os.ReadFile(fpPath)
 78  	if err != nil {
 79  		t.Fatalf("ReadFile %s: %v", fpPath, err)
 80  	}
 81  	if strings.TrimSpace(string(data)) != "name: foo" {
 82  		t.Errorf("unexpected content in %s: %q", fpPath, string(data))
 83  	}
 84  
 85  	// Verify vuln sub-directory file exists.
 86  	vulnPath := filepath.Join("data", "vuln", "bar", "CVE-2024-0001.yaml")
 87  	if _, err := os.Stat(vulnPath); err != nil {
 88  		t.Errorf("expected %s to exist: %v", vulnPath, err)
 89  	}
 90  
 91  	// Verify mcp was NOT copied (not in dirs list).
 92  	mcpPath := filepath.Join("data", "mcp", "tool.yaml")
 93  	if _, err := os.Stat(mcpPath); !os.IsNotExist(err) {
 94  		t.Errorf("expected %s to NOT exist", mcpPath)
 95  	}
 96  
 97  	// Verify README.md was NOT copied.
 98  	readmePath := filepath.Join(dstRoot, "README.md")
 99  	if _, err := os.Stat(readmePath); !os.IsNotExist(err) {
100  		t.Errorf("expected README.md to NOT be copied to dst root")
101  	}
102  }
103  
104  func TestCopyDataDirs_allDirs(t *testing.T) {
105  	srcRoot := t.TempDir()
106  	buildFakeClone(t, srcRoot)
107  
108  	dstRoot := t.TempDir()
109  	orig, _ := os.Getwd()
110  	if err := os.Chdir(dstRoot); err != nil {
111  		t.Fatalf("Chdir: %v", err)
112  	}
113  	defer os.Chdir(orig)
114  
115  	dirs := splitDirs(dataDirsDefault)
116  	n, err := copyDataDirs(srcRoot, dirs)
117  	if err != nil {
118  		t.Fatalf("copyDataDirs: %v", err)
119  	}
120  
121  	// fake clone has 3 data files: foo.yaml, CVE-2024-0001.yaml, tool.yaml
122  	if n != 3 {
123  		t.Errorf("expected 3 files written, got %d", n)
124  	}
125  }
126  
127  func TestCopyDataDirs_missingSubdir(t *testing.T) {
128  	srcRoot := t.TempDir()
129  	// Only create fingerprints, not vuln_en
130  	_ = os.MkdirAll(filepath.Join(srcRoot, "data", "fingerprints"), 0o755)
131  	_ = os.WriteFile(filepath.Join(srcRoot, "data", "fingerprints", "x.yaml"), []byte("x: 1\n"), 0o644)
132  
133  	dstRoot := t.TempDir()
134  	orig, _ := os.Getwd()
135  	if err := os.Chdir(dstRoot); err != nil {
136  		t.Fatalf("Chdir: %v", err)
137  	}
138  	defer os.Chdir(orig)
139  
140  	// vuln_en is missing in src — should be silently skipped, no error.
141  	dirs := []string{"fingerprints", "vuln_en"}
142  	n, err := copyDataDirs(srcRoot, dirs)
143  	if err != nil {
144  		t.Fatalf("expected no error for missing sub-dir, got: %v", err)
145  	}
146  	if n != 1 {
147  		t.Errorf("expected 1 file written, got %d", n)
148  	}
149  }
150  
151  func TestSplitDirs(t *testing.T) {
152  	cases := []struct {
153  		input string
154  		want  []string
155  	}{
156  		{"fingerprints,vuln", []string{"fingerprints", "vuln"}},
157  		{" fingerprints , vuln_en ", []string{"fingerprints", "vuln_en"}},
158  		{"", []string{}},
159  		{"mcp", []string{"mcp"}},
160  	}
161  	for _, tc := range cases {
162  		got := splitDirs(tc.input)
163  		if len(got) != len(tc.want) {
164  			t.Errorf("splitDirs(%q): got %v, want %v", tc.input, got, tc.want)
165  			continue
166  		}
167  		for i := range got {
168  			if got[i] != tc.want[i] {
169  				t.Errorf("splitDirs(%q)[%d]: got %q, want %q", tc.input, i, got[i], tc.want[i])
170  			}
171  		}
172  	}
173  }