Skip to content

Commit fbdff75

Browse files
gaohannkhopkiw
gaohannk
and
hopkiw
authored
add malformed ssh key unit test (#142)
* add malformed ssh key unit test and integration test * fix unit test name * fix go check * remove integration test and keep unit test * remove reflect deepEqual * refactor function removeExpiredKeys * rename * rename * go check * go check * fix * Apply suggestions from code review Co-authored-by: Liam Hopkins <[email protected]> * Apply suggestions from code review Co-authored-by: Liam Hopkins <[email protected]> * refactor * fix failed test and rename test * fix * address comments * fix Co-authored-by: Liam Hopkins <[email protected]>
1 parent b0c8cbd commit fbdff75

File tree

2 files changed

+90
-74
lines changed

2 files changed

+90
-74
lines changed

‎google_guest_agent/non_windows_accounts.go

Lines changed: 51 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ func (a *accountsMgr) diff() bool {
7373

7474
// If any on-disk keys have expired.
7575
for _, keys := range sshKeys {
76-
if len(keys) != len(removeExpiredKeys(keys)) {
76+
if len(keys) != len(getUserKeys(keys)) {
7777
return true
7878
}
7979
}
@@ -118,22 +118,7 @@ func (a *accountsMgr) set() error {
118118
mdkeys = append(mdkeys, newMetadata.Project.Attributes.SSHKeys...)
119119
}
120120

121-
mdKeyMap := make(map[string][]string)
122-
for _, key := range removeExpiredKeys(mdkeys) {
123-
idx := strings.Index(key, ":")
124-
if idx == -1 {
125-
logger.Debugf("invalid ssh key entry: %q", key)
126-
continue
127-
}
128-
user := key[:idx]
129-
if user == "" {
130-
logger.Debugf("invalid ssh key entry: %q", key)
131-
continue
132-
}
133-
userKeys := mdKeyMap[user]
134-
userKeys = append(userKeys, key[idx+1:])
135-
mdKeyMap[user] = userKeys
136-
}
121+
mdKeyMap := getUserKeys(mdkeys)
137122

138123
logger.Debugf("read google users file")
139124
gUsers, err := readGoogleUsersFile()
@@ -197,6 +182,55 @@ func (a *accountsMgr) set() error {
197182
return nil
198183
}
199184

185+
// getUserKeys returns the keys which are not expired and non-expiring key.
186+
// valid formats are:
187+
// user:ssh-rsa [KEY_VALUE] [USERNAME]
188+
// user:ssh-rsa [KEY_VALUE]
189+
// user:ssh-rsa [KEY_VALUE] google-ssh {"userName":"[USERNAME]","expireOn":"[EXPIRE_TIME]"}
190+
func getUserKeys(mdkeys []string) map[string][]string {
191+
mdKeyMap := make(map[string][]string)
192+
for i := 0; i < len(mdkeys); i++ {
193+
key := strings.Trim(mdkeys[i], " ")
194+
if key == "" {
195+
logger.Debugf("invalid ssh key entry: %q", key)
196+
continue
197+
}
198+
idx := strings.Index(key, ":")
199+
if idx == -1 {
200+
logger.Debugf("invalid ssh key entry: %q", key)
201+
continue
202+
}
203+
user := key[:idx]
204+
if user == "" {
205+
logger.Debugf("invalid ssh key entry: %q", key)
206+
continue
207+
}
208+
fields := strings.SplitN(key, " ", 4)
209+
if len(fields) == 3 && fields[2] == "google-ssh" {
210+
logger.Debugf("invalid ssh key entry: %q", key)
211+
// expiring key without expiration format.
212+
continue
213+
}
214+
if len(fields) > 3 {
215+
lkey := linuxKey{}
216+
if err := json.Unmarshal([]byte(fields[3]), &lkey); err != nil {
217+
// invalid expiration format.
218+
logger.Debugf("invalid ssh key entry: %q", key)
219+
continue
220+
}
221+
if lkey.expired() {
222+
logger.Debugf("expired ssh key entry: %q", key)
223+
continue
224+
}
225+
}
226+
// key which is not expired or non-expiring key, add it.
227+
userKeys := mdKeyMap[user]
228+
userKeys = append(userKeys, key[idx+1:])
229+
mdKeyMap[user] = userKeys
230+
}
231+
return mdKeyMap
232+
}
233+
200234
// passwdEntry is a user.User with omitted passwd fields restored.
201235
type passwdEntry struct {
202236
Username string
@@ -311,42 +345,6 @@ func (k linuxKey) expired() bool {
311345
return t.Before(time.Now())
312346
}
313347

314-
// removeExpiredKeys returns the provided list of keys with expired keys removed.
315-
// valid formats are:
316-
// ssh-rsa [KEY_VALUE] [USERNAME]
317-
// ssh-rsa [KEY_VALUE]
318-
// ssh-rsa [KEY_VALUE] google-ssh {"userName":"[USERNAME]","expireOn":"[EXPIRE_TIME]"}
319-
//
320-
// see: https://cloud.google.com/compute/docs/instances/adding-removing-ssh-keys#sshkeyformat
321-
func removeExpiredKeys(keys []string) []string {
322-
var res []string
323-
for i := 0; i < len(keys); i++ {
324-
key := strings.Trim(keys[i], " ")
325-
if key == "" {
326-
continue
327-
}
328-
fields := strings.SplitN(key, " ", 4)
329-
if len(fields) < 3 || fields[2] != "google-ssh" {
330-
// non-expiring key, add it.
331-
res = append(res, key)
332-
continue
333-
}
334-
if len(fields) < 4 {
335-
// expiring key without expiration format.
336-
continue
337-
}
338-
lkey := linuxKey{}
339-
if err := json.Unmarshal([]byte(fields[3]), &lkey); err != nil {
340-
// invalid expiration format.
341-
continue
342-
}
343-
if !lkey.expired() {
344-
res = append(res, key)
345-
}
346-
}
347-
return res
348-
}
349-
350348
// Replaces {user} or {group} in command string. Supports legacy python-era
351349
// user command overrides.
352350
func createUserGroupCmd(cmd, user, group string) *exec.Cmd {

‎google_guest_agent/windows_accounts_test.go

Lines changed: 39 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -216,32 +216,50 @@ func TestCompareAccounts(t *testing.T) {
216216
}
217217
}
218218

219-
func TestRemoveExpiredKeys(t *testing.T) {
219+
func TestGetUserKeys(t *testing.T) {
220220
var tests = []struct {
221-
key string
222-
valid bool
221+
key string
222+
expectedValid int
223223
}{
224-
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2028-11-08T19:30:47+0000"}`, true},
225-
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2028-11-08T19:30:47+0700"}`, true},
226-
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2028-11-08T19:30:47+0700", "futureField": "UNUSED_FIELDS_IGNORED"}`, true},
227-
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2018-11-08T19:30:46+0000"}`, false},
228-
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2018-11-08T19:30:46+0700"}`, false},
229-
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"INVALID_TIMESTAMP"}`, false},
230-
{`user:ssh-rsa [KEY] google-ssh`, false},
231-
{`user:ssh-rsa [KEY] user`, true},
232-
{`user:ssh-rsa [KEY]`, true},
233-
{},
224+
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2028-11-08T19:30:47+0000"}`,
225+
1,
226+
},
227+
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2028-11-08T19:30:47+0700"}`,
228+
1,
229+
},
230+
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2028-11-08T19:30:47+0700", "futureField": "UNUSED_FIELDS_IGNORED"}`,
231+
1,
232+
},
233+
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2018-11-08T19:30:46+0000"}`,
234+
0,
235+
},
236+
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"2018-11-08T19:30:46+0700"}`,
237+
0,
238+
},
239+
{`user:ssh-rsa [KEY] google-ssh {"userName":"[email protected]", "expireOn":"INVALID_TIMESTAMP"}`,
240+
0,
241+
},
242+
{`user:ssh-rsa [KEY] google-ssh`,
243+
0,
244+
},
245+
{`user:ssh-rsa [KEY] user`,
246+
1,
247+
},
248+
{`user:ssh-rsa [KEY]`,
249+
1,
250+
},
251+
{`malformed-ssh-keys [KEY] google-ssh`,
252+
0,
253+
},
254+
{`:malformed-ssh-keys [KEY] google-ssh`,
255+
0,
256+
},
234257
}
235258

236259
for _, tt := range tests {
237-
ret := removeExpiredKeys([]string{tt.key})
238-
if tt.valid {
239-
if len(ret) == 0 || ret[0] != tt.key {
240-
t.Errorf("valid key was removed: %q", tt.key)
241-
}
242-
}
243-
if !tt.valid && len(ret) == 1 {
244-
t.Errorf("invalid key was kept: %q", tt.key)
260+
ret := getUserKeys([]string{tt.key})
261+
if userKeys, _ := ret["user"]; len(userKeys) != tt.expectedValid {
262+
t.Errorf("expected %d valid keys from getUserKeys, but %d", tt.expectedValid, len(userKeys))
245263
}
246264
}
247265
}

0 commit comments

Comments
 (0)