Bladeren bron

normalized env variable names for bash

Darien Raymond 8 jaren geleden
bovenliggende
commit
b16a82024c
4 gewijzigde bestanden met toevoegingen van 96 en 19 verwijderingen
  1. 7 10
      common/buf/buffer_pool.go
  2. 42 0
      common/platform/platform.go
  3. 41 0
      common/platform/platform_test.go
  4. 6 9
      transport/ray/direct.go

+ 7 - 10
common/buf/buffer_pool.go

@@ -1,10 +1,10 @@
 package buf
 
 import (
-	"os"
 	"runtime"
-	"strconv"
 	"sync"
+
+	"v2ray.com/core/common/platform"
 )
 
 // Pool provides functionality to generate and recycle buffers on demand.
@@ -99,7 +99,7 @@ var (
 	mediumPool Pool
 )
 
-func getDefaultPoolSize() uint32 {
+func getDefaultPoolSize() int {
 	switch runtime.GOARCH {
 	case "amd64", "386":
 		return 20
@@ -109,14 +109,11 @@ func getDefaultPoolSize() uint32 {
 }
 
 func init() {
-	size := getDefaultPoolSize()
-	sizeStr := os.Getenv(poolSizeEnvKey)
-	if len(sizeStr) > 0 {
-		customSize, err := strconv.ParseUint(sizeStr, 10, 32)
-		if err == nil {
-			size = uint32(customSize)
-		}
+	f := platform.EnvFlag{
+		Name:    poolSizeEnvKey,
+		AltName: platform.NormalizeEnvName(poolSizeEnvKey),
 	}
+	size := f.GetValueAsInt(getDefaultPoolSize())
 	if size > 0 {
 		totalByteSize := size * 1024 * 1024
 		mediumPool = NewBufferPool(Size, totalByteSize/Size)

+ 42 - 0
common/platform/platform.go

@@ -0,0 +1,42 @@
+package platform
+
+import (
+	"os"
+	"strconv"
+	"strings"
+)
+
+type EnvFlag struct {
+	Name    string
+	AltName string
+}
+
+func (f EnvFlag) GetValue(defaultValue string) string {
+	if v, found := os.LookupEnv(f.Name); found {
+		return v
+	}
+	if len(f.AltName) > 0 {
+		if v, found := os.LookupEnv(f.AltName); found {
+			return v
+		}
+	}
+
+	return defaultValue
+}
+
+func (f EnvFlag) GetValueAsInt(defaultValue int) int {
+	const PlaceHolder = "xxxxxx"
+	s := f.GetValue(PlaceHolder)
+	if s == PlaceHolder {
+		return defaultValue
+	}
+	v, err := strconv.ParseInt(s, 10, 32)
+	if err != nil {
+		return defaultValue
+	}
+	return int(v)
+}
+
+func NormalizeEnvName(name string) string {
+	return strings.Replace(strings.ToUpper(strings.TrimSpace(name)), ".", "_", -1)
+}

+ 41 - 0
common/platform/platform_test.go

@@ -0,0 +1,41 @@
+package platform_test
+
+import (
+	"testing"
+
+	. "v2ray.com/core/common/platform"
+	"v2ray.com/core/testing/assert"
+)
+
+func TestNormalizeEnvName(t *testing.T) {
+	assert := assert.On(t)
+
+	cases := []struct {
+		input  string
+		output string
+	}{
+		{
+			input:  "a",
+			output: "A",
+		},
+		{
+			input:  "a.a",
+			output: "A_A",
+		},
+		{
+			input:  "A.A.B",
+			output: "A_A_B",
+		},
+	}
+	for _, test := range cases {
+		assert.String(NormalizeEnvName(test.input)).Equals(test.output)
+	}
+}
+
+func TestEnvFlag(t *testing.T) {
+	assert := assert.On(t)
+
+	assert.Int(EnvFlag{
+		Name: "xxxxx.y",
+	}.GetValueAsInt(10)).Equals(10)
+}

+ 6 - 9
transport/ray/direct.go

@@ -3,12 +3,11 @@ package ray
 import (
 	"context"
 	"io"
-	"os"
-	"strconv"
 	"sync"
 	"time"
 
 	"v2ray.com/core/common/buf"
+	"v2ray.com/core/common/platform"
 )
 
 // NewRay creates a new Ray for direct traffic transport.
@@ -44,13 +43,11 @@ var streamSizeLimit uint64 = 10 * 1024 * 1024
 
 func init() {
 	const raySizeEnvKey = "v2ray.ray.buffer.size"
-	sizeStr := os.Getenv(raySizeEnvKey)
-	if len(sizeStr) > 0 {
-		customSize, err := strconv.ParseUint(sizeStr, 10, 32)
-		if err == nil {
-			streamSizeLimit = customSize * 1024 * 1024
-		}
-	}
+	size := platform.EnvFlag{
+		Name:    raySizeEnvKey,
+		AltName: platform.NormalizeEnvName(raySizeEnvKey),
+	}.GetValueAsInt(10)
+	streamSizeLimit = uint64(size) * 1024 * 1024
 }
 
 type Stream struct {