Merge pull request #137 from Loyalsoldier/protoc-dont-rely-on-GOBIN-GOPATH

Protoc & vprotogen: do NOT rely on GOBIN & GOPATH

common/common.go

@@ -2,7 +2,16 @@
 // See each sub-package for detail.
 package common
 
-import "v2ray.com/core/common/errors"
+import (
+	"fmt"
+	"go/build"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"v2ray.com/core/common/errors"
+)
 
 //go:generate errorgen
 
@@ -28,3 +37,114 @@ func Must2(v interface{}, err error) interface{} {
 func Error2(v interface{}, err error) error {
 	return err
 }
+
+// envFile returns the name of the Go environment configuration file.
+// Copy from https://github.com/golang/go/blob/c4f2a9788a7be04daf931ac54382fbe2cb754938/src/cmd/go/internal/cfg/cfg.go#L150-L166
+func envFile() (string, error) {
+	if file := os.Getenv("GOENV"); file != "" {
+		if file == "off" {
+			return "", fmt.Errorf("GOENV=off")
+		}
+		return file, nil
+	}
+	dir, err := os.UserConfigDir()
+	if err != nil {
+		return "", err
+	}
+	if dir == "" {
+		return "", fmt.Errorf("missing user-config dir")
+	}
+	return filepath.Join(dir, "go", "env"), nil
+}
+
+// GetRuntimeEnv returns the value of runtime environment variable,
+// that is set by running following command: `go env -w key=value`.
+func GetRuntimeEnv(key string) (string, error) {
+	file, err := envFile()
+	if err != nil {
+		return "", err
+	}
+	if file == "" {
+		return "", fmt.Errorf("missing runtime env file")
+	}
+	var data []byte
+	var runtimeEnv string
+	data, readErr := ioutil.ReadFile(file)
+	if readErr != nil {
+		return "", readErr
+	}
+	envStrings := strings.Split(string(data), "\n")
+	for _, envItem := range envStrings {
+		envItem = strings.TrimSuffix(envItem, "\r")
+		envKeyValue := strings.Split(envItem, "=")
+		if strings.EqualFold(strings.TrimSpace(envKeyValue[0]), key) {
+			runtimeEnv = strings.TrimSpace(envKeyValue[1])
+		}
+	}
+	return runtimeEnv, nil
+}
+
+// GetGOBIN returns GOBIN environment variable as a string. It will NOT be empty.
+func GetGOBIN() string {
+	// The one set by user explicitly by `export GOBIN=/path` or `env GOBIN=/path command`
+	GOBIN := os.Getenv("GOBIN")
+	if GOBIN == "" {
+		var err error
+		// The one set by user by running `go env -w GOBIN=/path`
+		GOBIN, err = GetRuntimeEnv("GOBIN")
+		if err != nil {
+			// The default one that Golang uses
+			return filepath.Join(build.Default.GOPATH, "bin")
+		}
+		if GOBIN == "" {
+			return filepath.Join(build.Default.GOPATH, "bin")
+		}
+		return GOBIN
+	}
+	return GOBIN
+}
+
+// GetGOPATH returns GOPATH environment variable as a string. It will NOT be empty.
+func GetGOPATH() string {
+	// The one set by user explicitly by `export GOPATH=/path` or `env GOPATH=/path command`
+	GOPATH := os.Getenv("GOPATH")
+	if GOPATH == "" {
+		var err error
+		// The one set by user by running `go env -w GOPATH=/path`
+		GOPATH, err = GetRuntimeEnv("GOPATH")
+		if err != nil {
+			// The default one that Golang uses
+			return build.Default.GOPATH
+		}
+		if GOPATH == "" {
+			return build.Default.GOPATH
+		}
+		return GOPATH
+	}
+	return GOPATH
+}
+
+// GetModuleName returns the value of module in `go.mod` file.
+func GetModuleName(path string) (string, error) {
+	gomodPath := filepath.Join(path, "go.mod")
+	gomodBytes, err := ioutil.ReadFile(gomodPath)
+	if err != nil {
+		return "", err
+	}
+	gomodContent := string(gomodBytes)
+	moduleIdx := strings.Index(gomodContent, "module") + 6
+	newLineIdx := strings.Index(gomodContent, "\n")
+
+	var moduleName string
+	if moduleIdx >= 0 {
+		if newLineIdx >= 0 {
+			moduleName = strings.TrimSpace(gomodContent[moduleIdx:newLineIdx])
+			moduleName = strings.TrimSuffix(moduleName, "\r")
+		} else {
+			moduleName = strings.TrimSpace(gomodContent[moduleIdx:])
+		}
+	} else {
+		return "", fmt.Errorf("can not get the value of `module` in path `%s`", gomodPath)
+	}
+	return moduleName, nil
+}

infra/vprotogen/main.go

@@ -1,34 +1,36 @@
 package main
 
 import (
-	"flag"
 	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
 	"strings"
+
+	"v2ray.com/core/common"
 )
 
+var protoFilesUsingProtocGenGoFast = map[string]bool{"proxy/vless/encoding/addons.proto": true}
+
 var protocMap = map[string]string{
 	"windows": filepath.Join(".dev", "protoc", "windows", "protoc.exe"),
 	"darwin":  filepath.Join(".dev", "protoc", "macos", "protoc"),
 	"linux":   filepath.Join(".dev", "protoc", "linux", "protoc"),
 }
 
-var (
-	repo = flag.String("repo", "", "Repo for protobuf generation, such as v2ray.com/core")
-)
-
 func main() {
-	flag.Parse()
+	pwd, wdErr := os.Getwd()
+	if wdErr != nil {
+		fmt.Println("Can not get current working directory.")
+		os.Exit(1)
+	}
 
-	protofiles := make(map[string][]string)
+	GOBIN := common.GetGOBIN()
 	protoc := protocMap[runtime.GOOS]
-	gosrc := filepath.Join(os.Getenv("GOPATH"), "src")
-	reporoot := filepath.Join(os.Getenv("GOPATH"), "src", *repo)
 
-	filepath.Walk(reporoot, func(path string, info os.FileInfo, err error) error {
+	protoFilesMap := make(map[string][]string)
+	walkErr := filepath.Walk("./", func(path string, info os.FileInfo, err error) error {
 		if err != nil {
 			fmt.Println(err)
 			return err
@@ -41,33 +43,114 @@ func main() {
 		dir := filepath.Dir(path)
 		filename := filepath.Base(path)
 		if strings.HasSuffix(filename, ".proto") {
-			protofiles[dir] = append(protofiles[dir], path)
+			protoFilesMap[dir] = append(protoFilesMap[dir], path)
 		}
 
 		return nil
 	})
-
-	var protoFilesUsingProtocGenGoFast = map[string]bool{"proxy/vless/encoding/addons.proto": true}
-
-	for _, files := range protofiles {
-		for _, absPath := range files {
-			relPath, _ := filepath.Rel(reporoot, absPath)
-			args := make([]string, 0)
-			if protoFilesUsingProtocGenGoFast[relPath] {
-				args = []string{"--proto_path", reporoot, "--gofast_out", gosrc}
-			} else {
-				args = []string{"--proto_path", reporoot, "--go_out", gosrc, "--go-grpc_out", gosrc}
+	if walkErr != nil {
+		fmt.Println(walkErr)
+		os.Exit(1)
 	}
-			args = append(args, absPath)
+
+	for _, files := range protoFilesMap {
+		for _, relProtoFile := range files {
+			var args []string
+			if protoFilesUsingProtocGenGoFast[relProtoFile] {
+				args = []string{"--gofast_out", pwd, "--plugin", "protoc-gen-gofast=" + GOBIN + "/protoc-gen-gofast"}
+			} else {
+				args = []string{"--go_out", pwd, "--go-grpc_out", pwd, "--plugin", "protoc-gen-go=" + GOBIN + "/protoc-gen-go", "--plugin", "protoc-gen-go-grpc=" + GOBIN + "/protoc-gen-go-grpc"}
+			}
+			args = append(args, relProtoFile)
 			cmd := exec.Command(protoc, args...)
 			cmd.Env = append(cmd.Env, os.Environ()...)
-			output, err := cmd.CombinedOutput()
+			cmd.Env = append(cmd.Env, "GOBIN="+GOBIN)
+			output, cmdErr := cmd.CombinedOutput()
 			if len(output) > 0 {
 				fmt.Println(string(output))
 			}
+			if cmdErr != nil {
+				fmt.Println(cmdErr)
+				os.Exit(1)
+			}
+		}
+	}
+
+	moduleName, gmnErr := common.GetModuleName(pwd)
+	if gmnErr != nil {
+		fmt.Println(gmnErr)
+		os.Exit(1)
+	}
+	modulePath := filepath.Join(strings.Split(moduleName, string(os.PathSeparator))...)
+
+	pbGoFilesMap := make(map[string][]string)
+	walkErr2 := filepath.Walk(modulePath, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			fmt.Println(err)
+			return err
+		}
+
+		if info.IsDir() {
+			return nil
+		}
+
+		dir := filepath.Dir(path)
+		filename := filepath.Base(path)
+		if strings.HasSuffix(filename, ".pb.go") {
+			pbGoFilesMap[dir] = append(pbGoFilesMap[dir], path)
+		}
+
+		return nil
+	})
+	if walkErr2 != nil {
+		fmt.Println(walkErr2)
+		os.Exit(1)
+	}
+
+	var err error
+	for _, srcPbGoFiles := range pbGoFilesMap {
+		for _, srcPbGoFile := range srcPbGoFiles {
+			var dstPbGoFile string
+			dstPbGoFile, err = filepath.Rel(modulePath, srcPbGoFile)
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			err = os.Link(srcPbGoFile, dstPbGoFile)
+			if err != nil {
+				if os.IsNotExist(err) {
+					fmt.Printf("'%s' does not exist\n", srcPbGoFile)
+					continue
+				}
+				if os.IsPermission(err) {
+					fmt.Println(err)
+					continue
+				}
+				if os.IsExist(err) {
+					err = os.Remove(dstPbGoFile)
+					if err != nil {
+						fmt.Printf("Failed to delete file '%s'\n", dstPbGoFile)
+						continue
+					}
+					err = os.Rename(srcPbGoFile, dstPbGoFile)
+					if err != nil {
+						fmt.Printf("Can not move '%s' to '%s'\n", srcPbGoFile, dstPbGoFile)
+					}
+					continue
+				}
+			}
+			err = os.Rename(srcPbGoFile, dstPbGoFile)
+			if err != nil {
+				fmt.Printf("Can not move '%s' to '%s'\n", srcPbGoFile, dstPbGoFile)
+			}
+			continue
+		}
+	}
+
+	if err == nil {
+		err = os.RemoveAll(strings.Split(modulePath, string(os.PathSeparator))[0])
 		if err != nil {
 			fmt.Println(err)
 		}
 	}
 }
-}

proto.go

@@ -1,8 +1,7 @@
 package core
 
-//go:generate go install "google.golang.org/protobuf/proto"
-//go:generate go install "google.golang.org/protobuf/cmd/protoc-gen-go"
-//go:generate go install "google.golang.org/grpc/cmd/protoc-gen-go-grpc"
-//go:generate go install "github.com/gogo/protobuf/protoc-gen-gofast"
-//go:generate go install "v2ray.com/core/infra/vprotogen"
-//go:generate vprotogen -repo v2ray.com/core
+//go:generate go install google.golang.org/protobuf/proto
+//go:generate go install google.golang.org/protobuf/cmd/protoc-gen-go
+//go:generate go get -v google.golang.org/grpc/cmd/protoc-gen-go-grpc@v0.0.0-20200825170228-39ef2aaf62df
+//go:generate go install github.com/gogo/protobuf/protoc-gen-gofast
+//go:generate go run v2ray.com/core/infra/vprotogen