From 025ae38acc409137543a843742fae8d58ad07d8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BE=90=E8=83=96?= <xuyang2436@gmail.com>
Date: Mon, 26 Jul 2021 20:08:48 +0800
Subject: [PATCH] fix(config/env): fix env.load() index out of range (#1252)

* fix(config/env): env.load() may panic
---
 config/env/env.go      | 12 +++---
 config/env/env_test.go | 95 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 7 deletions(-)

diff --git a/config/env/env.go b/config/env/env.go
index c7d018b63..b024fc0a8 100644
--- a/config/env/env.go
+++ b/config/env/env.go
@@ -30,15 +30,13 @@ func (e *env) load(envStrings []string) []*config.KeyValue {
 		}
 
 		if len(e.prefixs) > 0 {
-			p, ok := matchPrefix(e.prefixs, envstr)
+			p, ok := matchPrefix(e.prefixs, k)
 			if !ok || len(p) == len(k) {
 				continue
 			}
 			// trim prefix
-			k = k[len(p):]
-			if k[0] == '_' {
-				k = k[1:]
-			}
+			k = strings.TrimPrefix(k, p)
+			k = strings.TrimPrefix(k, "_")
 		}
 
 		if len(k) != 0 {
@@ -59,8 +57,8 @@ func (e *env) Watch() (config.Watcher, error) {
 	return w, nil
 }
 
-func matchPrefix(prefixs []string, s string) (string, bool) {
-	for _, p := range prefixs {
+func matchPrefix(prefixes []string, s string) (string, bool) {
+	for _, p := range prefixes {
 		if strings.HasPrefix(s, p) {
 			return p, true
 		}
diff --git a/config/env/env_test.go b/config/env/env_test.go
index 13861da8c..0e19577b1 100644
--- a/config/env/env_test.go
+++ b/config/env/env_test.go
@@ -281,6 +281,44 @@ func Test_env_load(t *testing.T) {
 			},
 		},
 
+		{
+			name: "empty prefix",
+			fields: fields{
+				prefixs: []string{""},
+			},
+			args: args{
+				envStrings: []string{
+					"__SERVICE_NAME=kratos_app",
+					"__ADDR=192.168.0.1",
+					"__AGE=20",
+				},
+			},
+			want: []*config.KeyValue{
+				{Key: "_SERVICE_NAME", Value: []byte("kratos_app"), Format: ""},
+				{Key: "_ADDR", Value: []byte("192.168.0.1"), Format: ""},
+				{Key: "_AGE", Value: []byte("20"), Format: ""},
+			},
+		},
+
+		{
+			name: "underscore prefix",
+			fields: fields{
+				prefixs: []string{"_"},
+			},
+			args: args{
+				envStrings: []string{
+					"__SERVICE_NAME=kratos_app",
+					"__ADDR=192.168.0.1",
+					"__AGE=20",
+				},
+			},
+			want: []*config.KeyValue{
+				{Key: "SERVICE_NAME", Value: []byte("kratos_app"), Format: ""},
+				{Key: "ADDR", Value: []byte("192.168.0.1"), Format: ""},
+				{Key: "AGE", Value: []byte("20"), Format: ""},
+			},
+		},
+
 		{
 			name: "with prefixes",
 			fields: fields{
@@ -299,6 +337,32 @@ func Test_env_load(t *testing.T) {
 				{Key: "AGE", Value: []byte("20"), Format: ""},
 			},
 		},
+
+		{
+			name: "should not panic #1",
+			fields: fields{
+				prefixs: []string{"FOO"},
+			},
+			args: args{
+				envStrings: []string{
+					"FOO=123",
+				},
+			},
+			want: nil,
+		},
+
+		{
+			name: "should not panic #2",
+			fields: fields{
+				prefixs: []string{"FOO=1"},
+			},
+			args: args{
+				envStrings: []string{
+					"FOO=123",
+				},
+			},
+			want: nil,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -312,3 +376,34 @@ func Test_env_load(t *testing.T) {
 		})
 	}
 }
+
+func Test_matchPrefix(t *testing.T) {
+	type args struct {
+		prefixes []string
+		s        string
+	}
+	tests := []struct {
+		name   string
+		args   args
+		want   string
+		wantOk bool
+	}{
+		{args: args{prefixes: nil, s: "foo=123"}, want: "", wantOk: false},
+		{args: args{prefixes: []string{""}, s: "foo=123"}, want: "", wantOk: true},
+		{args: args{prefixes: []string{"foo"}, s: "foo=123"}, want: "foo", wantOk: true},
+		{args: args{prefixes: []string{"foo=1"}, s: "foo=123"}, want: "foo=1", wantOk: true},
+		{args: args{prefixes: []string{"foo=1234"}, s: "foo=123"}, want: "", wantOk: false},
+		{args: args{prefixes: []string{"bar"}, s: "foo=123"}, want: "", wantOk: false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, gotOk := matchPrefix(tt.args.prefixes, tt.args.s)
+			if got != tt.want {
+				t.Errorf("matchPrefix() got = %v, want %v", got, tt.want)
+			}
+			if gotOk != tt.wantOk {
+				t.Errorf("matchPrefix() gotOk = %v, wantOk %v", gotOk, tt.wantOk)
+			}
+		})
+	}
+}