diff --git a/.envrc b/.envrc
new file mode 100644
index 0000000..3550a30
--- /dev/null
+++ b/.envrc
@@ -0,0 +1 @@
+use flake
diff --git a/.gitignore b/.gitignore
index 5c634a1..4578cd7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,3 +18,4 @@
 .DS_Store
 *.wasm
 dist/
+.direnv/
\ No newline at end of file
diff --git a/cmd/protoc-gen-go-plugin/main.go b/cmd/protoc-gen-go-plugin/main.go
index 6945787..23b5aef 100644
--- a/cmd/protoc-gen-go-plugin/main.go
+++ b/cmd/protoc-gen-go-plugin/main.go
@@ -12,8 +12,15 @@ import (
 func main() {
 	var flags flag.FlagSet
 	disablePbGen := flags.Bool("disable_pb_gen", false, "disable .pb.go generation")
+	wasmPackage := flags.String("wasm_package", "github.com/knqyf263/go-plugin/wasm", "override package that provide wasm memory management")
+	useGoPluginKnownTypes := flags.Bool("goplugin_known_types", true, "use go-plugin known types")
 	protogen.Options{ParamFunc: flags.Set}.Run(func(plugin *protogen.Plugin) error {
-		g, err := gen.NewGenerator(plugin)
+		opts := gen.Options{
+			UseGoPluginKnownTypes: *useGoPluginKnownTypes,
+			DisablePBGen:          *disablePbGen,
+			WasmPackage:           *wasmPackage,
+		}
+		g, err := gen.NewGenerator(plugin, opts)
 		if err != nil {
 			return err
 		}
@@ -22,7 +29,7 @@ func main() {
 			if !f.Generate {
 				continue
 			}
-			g.GenerateFiles(f, gen.Options{DisablePBGen: *disablePbGen})
+			g.GenerateFiles(f, opts)
 		}
 
 		plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
diff --git a/flake.lock b/flake.lock
new file mode 100644
index 0000000..4fc8f58
--- /dev/null
+++ b/flake.lock
@@ -0,0 +1,61 @@
+{
+  "nodes": {
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1761597516,
+        "narHash": "sha256-wxX7u6D2rpkJLWkZ2E932SIvDJW8+ON/0Yy8+a5vsDU=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "daf6dc47aa4b44791372d6139ab7b25269184d55",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-25.05",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}
diff --git a/flake.nix b/flake.nix
new file mode 100644
index 0000000..9ddfdd1
--- /dev/null
+++ b/flake.nix
@@ -0,0 +1,29 @@
+{
+  description = "bitmagnet dev shell";
+  inputs.nixpkgs.url = "github:nixos/nixpkgs/nixos-25.05";
+  inputs.flake-utils.url = "github:numtide/flake-utils";
+
+  outputs = {
+    nixpkgs,
+    flake-utils,
+    ...
+  }:
+    flake-utils.lib.eachDefaultSystem (system: let
+    pkgs = import nixpkgs {
+      system = system;
+    };
+    in {
+      formatter = pkgs.alejandra;
+      devShells = {
+        default = pkgs.mkShell {
+          hardeningDisable = [ "fortify" ]; 
+          packages = with pkgs; [
+            go
+            golangci-lint
+            protobuf
+            protoc-gen-go
+          ];
+        };
+      };
+    });
+}
diff --git a/gen/main.go b/gen/main.go
index 05f1430..b848df7 100644
--- a/gen/main.go
+++ b/gen/main.go
@@ -76,7 +76,7 @@ var (
 	wazeroSysPackage  goImportPath = protogen.GoImportPath("github.com/tetratelabs/wazero/sys")
 	wazeroWasiPackage goImportPath = protogen.GoImportPath("github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1")
 
-	pluginWasmPackage goImportPath = protogen.GoImportPath("github.com/knqyf263/go-plugin/wasm")
+	pluginWasmPackage goImportPath = protogen.GoImportPath("")
 )
 
 type goImportPath interface {
@@ -90,10 +90,12 @@ type Generator struct {
 }
 
 type Options struct {
-	DisablePBGen bool
+	UseGoPluginKnownTypes bool
+	DisablePBGen          bool
+	WasmPackage           string
 }
 
-func NewGenerator(plugin *protogen.Plugin) (*Generator, error) {
+func NewGenerator(plugin *protogen.Plugin, opts Options) (*Generator, error) {
 	ext := &vtgenerator.Extensions{}
 	featureNames := []string{"marshal", "unmarshal", "size"}
 
@@ -103,7 +105,7 @@ func NewGenerator(plugin *protogen.Plugin) (*Generator, error) {
 	}
 
 	for _, f := range plugin.Files {
-		if !f.Generate {
+		if !f.Generate || !opts.UseGoPluginKnownTypes {
 			continue
 		}
 
@@ -144,6 +146,7 @@ func replaceImport(m *protogen.Message) {
 // GenerateFiles generates the contents of a .pb.go file.
 func (gg *Generator) GenerateFiles(file *protogen.File, opts Options) *protogen.GeneratedFile {
 	f := gg.newFileInfo(file)
+	pluginWasmPackage = protogen.GoImportPath(opts.WasmPackage)
 	gg.generatePBFile(f, opts.DisablePBGen)
 	gg.generateHostFile(f)
 	gg.generatePluginFile(f)