1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
package cctx
import (
"context"
"github.com/gliderlabs/ssh"
"github.com/jumpserver/koko/pkg/model"
)
type contextKey struct {
name string
}
var (
ContextKeyUser = &contextKey{"user"}
ContextKeyAsset = &contextKey{"asset"}
ContextKeySystemUser = &contextKey{"systemUser"}
ContextKeySSHSession = &contextKey{"sshSession"}
ContextKeyLocalAddr = &contextKey{"localAddr"}
ContextKeyRemoteAddr = &contextKey{"RemoteAddr"}
ContextKeySSHCtx = &contextKey{"sshCtx"}
ContextKeySeed = &contextKey{"seed"}
ContextKeyToken = &contextKey{"token"}
)
type Context interface {
context.Context
User() *model.User
Asset() *model.Asset
SystemUser() *model.SystemUser
SSHSession() *ssh.Session
SSHCtx() *ssh.Context
SetValue(key, value interface{})
}
// Context coco内部使用的Context
type CocoContext struct {
context.Context
}
// user 返回当前连接的用户model
func (ctx *CocoContext) User() *model.User {
return ctx.Value(ContextKeyUser).(*model.User)
}
func (ctx *CocoContext) Asset() *model.Asset {
return ctx.Value(ContextKeyAsset).(*model.Asset)
}
func (ctx *CocoContext) SystemUser() *model.SystemUser {
return ctx.Value(ContextKeySystemUser).(*model.SystemUser)
}
func (ctx *CocoContext) SSHSession() *ssh.Session {
return ctx.Value(ContextKeySSHSession).(*ssh.Session)
}
func (ctx *CocoContext) SSHCtx() *ssh.Context {
return ctx.Value(ContextKeySSHCtx).(*ssh.Context)
}
func (ctx *CocoContext) SetValue(key, value interface{}) {
ctx.Context = context.WithValue(ctx.Context, key, value)
}
func applySessionMetadata(ctx *CocoContext, sess ssh.Session) {
ctx.SetValue(ContextKeySSHSession, &sess)
ctx.SetValue(ContextKeySSHCtx, sess.Context())
ctx.SetValue(ContextKeyLocalAddr, sess.LocalAddr())
}
func NewContext(sess ssh.Session) (*CocoContext, context.CancelFunc) {
sshCtx, cancel := context.WithCancel(sess.Context())
ctx := &CocoContext{sshCtx}
applySessionMetadata(ctx, sess)
return ctx, cancel
}