feat(capi+go): implement API for banning the use of certain modules.

Author: plusvic

capi/include/yara_x.h

@@ -260,6 +260,18 @@ enum YRX_RESULT yrx_compiler_add_source_with_origin(struct YRX_COMPILER *compile
 enum YRX_RESULT yrx_compiler_ignore_module(struct YRX_COMPILER *compiler,
                                            const char *module);
 
+// Tell the compiler that a YARA module can't be used.
+//
+// Import statements for the banned module will cause an error. The error
+// message can be customized by using the given error title and message.
+//
+// If this function is called multiple times with the same module name,
+// the error title and message will be updated.
+enum YRX_RESULT yrx_compiler_ban_module(struct YRX_COMPILER *compiler,
+                                        const char *module,
+                                        const char *error_title,
+                                        const char *error_msg);
+
 // Creates a new namespace.
 //
 // Further calls to `yrx_compiler_add_source` will put the rules under the

capi/src/compiler.rs

@@ -159,6 +159,50 @@ pub unsafe extern "C" fn yrx_compiler_ignore_module(
     YRX_RESULT::SUCCESS
 }
 
+/// Tell the compiler that a YARA module can't be used.
+///
+/// Import statements for the banned module will cause an error. The error
+/// message can be customized by using the given error title and message.
+///
+/// If this function is called multiple times with the same module name,
+/// the error title and message will be updated.
+#[no_mangle]
+pub unsafe extern "C" fn yrx_compiler_ban_module(
+    compiler: *mut YRX_COMPILER,
+    module: *const c_char,
+    error_title: *const c_char,
+    error_msg: *const c_char,
+) -> YRX_RESULT {
+    let compiler = if let Some(compiler) = compiler.as_mut() {
+        compiler
+    } else {
+        return YRX_RESULT::INVALID_ARGUMENT;
+    };
+
+    let module = if let Ok(module) = CStr::from_ptr(module).to_str() {
+        module
+    } else {
+        return YRX_RESULT::INVALID_ARGUMENT;
+    };
+
+    let err_title = if let Ok(err_title) = CStr::from_ptr(error_title).to_str()
+    {
+        err_title
+    } else {
+        return YRX_RESULT::INVALID_ARGUMENT;
+    };
+
+    let err_msg = if let Ok(err_msg) = CStr::from_ptr(error_msg).to_str() {
+        err_msg
+    } else {
+        return YRX_RESULT::INVALID_ARGUMENT;
+    };
+
+    compiler.inner.ban_module(module, err_title, err_msg);
+
+    YRX_RESULT::SUCCESS
+}
+
 /// Creates a new namespace.
 ///
 /// Further calls to `yrx_compiler_add_source` will put the rules under the

go/compiler.go

@@ -51,6 +51,21 @@ func IgnoreModule(module string) CompileOption {
 	}
 }
 
+// BanModule is an option for [NewCompiler] and [Compile] that allows
+// banning the use of a given module.
+//
+// Import statements for the banned module will cause an error. The error
+// message can be customized by using the given error title and message.
+//
+// If this function is called multiple times with the same module name,
+// the error title and message will be updated.
+func BanModule(module string, errTitle string, errMessage string) CompileOption {
+	return func(c *Compiler) error {
+		c.bannedModules[module] = bannedModule{errTitle, errMessage}
+		return nil
+	}
+}
+
 // RelaxedReSyntax is an option for [NewCompiler] and [Compile] that
 // determines whether the compiler should adopt a more relaxed approach
 // while parsing regular expressions.
@@ -104,7 +119,7 @@ type SourceOption func(opt *sourceOptions) error
 // The origin is usually the path of the file containing the source code,
 // but it can be any arbitrary string that conveys information of the
 // source's origin. This origin appears in error reports, for instance, if
-// if origin is "some_file.yar", error reports will look like:
+// origin is "some_file.yar", error reports will look like:
 //
 //	error: syntax error
 //	 --> some_file.yar:4:17
@@ -194,20 +209,27 @@ func (c CompileError) Error() string {
 	return c.Text
 }
 
+type bannedModule struct {
+	errTitle string
+	errMsg   string
+}
+
 // Compiler represent a YARA compiler.
 type Compiler struct {
 	cCompiler          *C.YRX_COMPILER
 	relaxedReSyntax    bool
 	errorOnSlowPattern bool
 	errorOnSlowLoop    bool
 	ignoredModules     map[string]bool
+	bannedModules      map[string]bannedModule
 	vars               map[string]interface{}
 }
 
 // NewCompiler creates a new compiler.
 func NewCompiler(opts ...CompileOption) (*Compiler, error) {
 	c := &Compiler{
 		ignoredModules: make(map[string]bool),
+		bannedModules:  make(map[string]bannedModule),
 		vars:           make(map[string]interface{}),
 	}
 
@@ -244,6 +266,9 @@ func (c *Compiler) initialize() error {
 	for name, _ := range c.ignoredModules {
 		c.ignoreModule(name)
 	}
+	for name, v := range c.bannedModules {
+		c.banModule(name, v.errTitle, v.errMsg)
+	}
 	for ident, value := range c.vars {
 		if err := c.DefineGlobal(ident, value); err != nil {
 			return err
@@ -325,6 +350,23 @@ func (c *Compiler) ignoreModule(module string) {
 	runtime.KeepAlive(c)
 }
 
+func (c *Compiler) banModule(module, error_title, error_message string) {
+	cModule := C.CString(module)
+	defer C.free(unsafe.Pointer(cModule))
+
+	cErrTitle := C.CString(error_title)
+	defer C.free(unsafe.Pointer(cErrTitle))
+
+	cErrMsg := C.CString(error_message)
+	defer C.free(unsafe.Pointer(cErrMsg))
+
+	result := C.yrx_compiler_ban_module(c.cCompiler, cModule, cErrTitle, cErrMsg)
+	if result != C.SUCCESS {
+		panic("yrx_compiler_add_unsupported_module failed")
+	}
+	runtime.KeepAlive(c)
+}
+
 // NewNamespace creates a new namespace.
 //
 // Later calls to [Compiler.AddSource] will put the rules under the newly created

go/compiler_test.go

@@ -31,6 +31,20 @@ func TestUnsupportedModules(t *testing.T) {
 	assert.Len(t, scanResults.MatchingRules(), 1)
 }
 
+func TestBannedModules(t *testing.T) {
+	_, err := Compile(
+		`import "pe"`,
+		BanModule("pe", "pe module is banned", "pe module was used here"))
+
+	expected := `error[E100]: pe module is banned
+ --> line:1:1
+  |
+1 | import "pe"
+  | ^^^^^^^^^^^ pe module was used here
+  |`
+	assert.EqualError(t, err, expected)
+}
+
 func TestRelaxedReSyntax(t *testing.T) {
 	r, err := Compile(`
 		rule test { strings: $a = /\Release/ condition: $a }`,