-
Notifications
You must be signed in to change notification settings - Fork 217
/
grpc_handler.go
150 lines (138 loc) · 3.86 KB
/
grpc_handler.go
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
package auth
import (
"bufio"
"container/list"
"context"
"io/ioutil"
"os"
"github.com/golang/protobuf/ptypes/empty"
"go.uber.org/zap"
"gopkg.in/yaml.v2"
"github.com/DrmagicE/gmqtt/plugin/admin"
)
// List lists all accounts
func (a *Auth) List(ctx context.Context, req *ListAccountsRequest) (resp *ListAccountsResponse, err error) {
page, pageSize := admin.GetPage(req.Page, req.PageSize)
offset, n := admin.GetOffsetN(page, pageSize)
a.mu.RLock()
defer a.mu.RUnlock()
resp = &ListAccountsResponse{
Accounts: []*Account{},
TotalCount: 0,
}
a.indexer.Iterate(func(elem *list.Element) {
resp.Accounts = append(resp.Accounts, elem.Value.(*Account))
}, offset, n)
resp.TotalCount = uint32(a.indexer.Len())
return resp, nil
}
// Get gets the account for given username.
// Return NotFound error when account not found.
func (a *Auth) Get(ctx context.Context, req *GetAccountRequest) (resp *GetAccountResponse, err error) {
if req.Username == "" {
return nil, admin.ErrInvalidArgument("username", "cannot be empty")
}
a.mu.RLock()
defer a.mu.RUnlock()
resp = &GetAccountResponse{}
if e := a.indexer.GetByID(req.Username); e != nil {
resp.Account = e.Value.(*Account)
return resp, nil
}
return nil, admin.ErrNotFound
}
// saveFileHandler is the default handler for auth.saveFile, must call after auth.mu is locked
func (a *Auth) saveFileHandler() error {
tmpfile, err := ioutil.TempFile("./", "gmqtt_password")
if err != nil {
return err
}
w := bufio.NewWriter(tmpfile)
// get all accounts
var accounts []*Account
a.indexer.Iterate(func(elem *list.Element) {
accounts = append(accounts, elem.Value.(*Account))
}, 0, uint(a.indexer.Len()))
b, err := yaml.Marshal(accounts)
if err != nil {
return err
}
_, err = w.Write(b)
if err != nil {
return err
}
err = w.Flush()
if err != nil {
return err
}
tmpfile.Close()
// replace the old password file.
return os.Rename(tmpfile.Name(), a.config.PasswordFile)
}
// Update updates the password for the account.
// Create a new account if the account for the username is not exists.
// Update will persist the account data to the password file.
func (a *Auth) Update(ctx context.Context, req *UpdateAccountRequest) (resp *empty.Empty, err error) {
if req.Username == "" {
return nil, admin.ErrInvalidArgument("username", "cannot be empty")
}
hashedPassword, err := a.generatePassword(req.Password)
if err != nil {
return &empty.Empty{}, err
}
a.mu.Lock()
defer a.mu.Unlock()
var oact *Account
elem := a.indexer.GetByID(req.Username)
if elem != nil {
oact = elem.Value.(*Account)
}
a.indexer.Set(req.Username, &Account{
Username: req.Username,
Password: hashedPassword,
})
err = a.saveFile()
if err != nil {
// should rollback if failed to persist to file.
if oact == nil {
a.indexer.Remove(req.Username)
return &empty.Empty{}, err
}
a.indexer.Set(req.Username, &Account{
Username: req.Username,
Password: oact.Password,
})
}
if oact == nil {
log.Info("new account created", zap.String("username", req.Username))
} else {
log.Info("password updated", zap.String("username", req.Username))
}
return &empty.Empty{}, err
}
// Delete deletes the account for the username.
func (a *Auth) Delete(ctx context.Context, req *DeleteAccountRequest) (resp *empty.Empty, err error) {
if req.Username == "" {
return nil, admin.ErrInvalidArgument("username", "cannot be empty")
}
a.mu.Lock()
defer a.mu.Unlock()
act := a.indexer.GetByID(req.Username)
if act == nil {
// fast path
return &empty.Empty{}, nil
}
oact := act.Value
a.indexer.Remove(req.Username)
err = a.saveFile()
if err != nil {
// should rollback if failed to persist to file
a.indexer.Set(req.Username, &Account{
Username: req.Username,
Password: oact.(*Account).Password,
})
return &empty.Empty{}, err
}
log.Info("account deleted", zap.String("username", req.Username))
return &empty.Empty{}, nil
}