-
Notifications
You must be signed in to change notification settings - Fork 51
/
cgnetcls.go
303 lines (236 loc) · 8.08 KB
/
cgnetcls.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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
// +build linux,!darwin,!windows
//Package cgnetcls implements functionality to manage classid for processes belonging to different cgroups
package cgnetcls
import (
"bufio"
"errors"
"fmt"
"io/ioutil"
"os"
"path/filepath"
"strconv"
"strings"
"sync/atomic"
"syscall"
"github.com/kardianos/osext"
"go.uber.org/zap"
)
//Initialize only ince
func init() {
mountCgroupController()
}
// Creategroup creates a cgroup/net_cls structure and writes the allocated classid to the file.
// To add a new process to this cgroup we need to write to the cgroup file
func (s *netCls) Creategroup(cgroupname string) error {
//Create the directory structure
_, err := os.Stat(basePath + procs)
if os.IsNotExist(err) {
if err = syscall.Mount("cgroup", basePath, "cgroup", 0, "net_cls,net_prio"); err != nil {
return err
}
}
cgroupPath := filepath.Join(basePath, s.TriremePath, cgroupname)
if _, err = os.Stat(cgroupPath); err == nil {
return nil
}
if err = os.MkdirAll(cgroupPath, 0700); err != nil {
return err
}
// Write to the notify on release file and release agent files
if s.ReleaseAgentPath != "" {
err = ioutil.WriteFile(filepath.Join(basePath, releaseAgentConfFile), []byte(s.ReleaseAgentPath), 0644)
if err != nil {
return fmt.Errorf("unable to register a release agent error: %s", err)
}
err = ioutil.WriteFile(filepath.Join(basePath, notifyOnReleaseFile), []byte("1"), 0644)
if err != nil {
return fmt.Errorf("unable to write to the notify file: %s", err)
}
err = ioutil.WriteFile(filepath.Join(basePath, s.TriremePath, notifyOnReleaseFile), []byte("1"), 0644)
if err != nil {
return fmt.Errorf("unable to write to the notify file: %s", err)
}
err = ioutil.WriteFile(filepath.Join(basePath, s.TriremePath, cgroupname, notifyOnReleaseFile), []byte("1"), 0644)
if err != nil {
return fmt.Errorf("unable to write to the notify file: %s", err)
}
}
return nil
}
//AssignMark writes the mark value to net_cls.classid file.
func (s *netCls) AssignMark(cgroupname string, mark uint64) error {
_, err := os.Stat(filepath.Join(basePath, s.TriremePath, cgroupname))
if os.IsNotExist(err) {
return fmt.Errorf("cgroup does not exist: %s", err)
}
//16 is the base since the mark file expects hexadecimal values
markval := "0x" + (strconv.FormatUint(mark, 16))
if err := ioutil.WriteFile(filepath.Join(basePath, s.TriremePath, cgroupname, markFile), []byte(markval), 0644); err != nil {
return fmt.Errorf("failed to write to net_cls.classid file for new cgroup: %s", err)
}
return nil
}
// AddProcess adds the process to the net_cls group
func (s *netCls) AddProcess(cgroupname string, pid int) error {
_, err := os.Stat(filepath.Join(basePath, s.TriremePath, cgroupname))
if os.IsNotExist(err) {
return fmt.Errorf("cannot add process. cgroup does not exist: %s", err)
}
PID := []byte(strconv.Itoa(pid))
if err := syscall.Kill(pid, 0); err != nil {
return nil
}
if err := ioutil.WriteFile(filepath.Join(basePath, s.TriremePath, cgroupname, procs), PID, 0644); err != nil {
return fmt.Errorf("cannot add process: %s", err)
}
return nil
}
//RemoveProcess removes the process from the cgroup by writing the pid to the
//top of net_cls cgroup cgroup.procs
func (s *netCls) RemoveProcess(cgroupname string, pid int) error {
_, err := os.Stat(filepath.Join(basePath, s.TriremePath, cgroupname))
if os.IsNotExist(err) {
return fmt.Errorf("cannot clean up process. cgroup does not exist: %s", err)
}
data, err := ioutil.ReadFile(filepath.Join(basePath, procs))
if err != nil {
return fmt.Errorf("cannot cleanup process: %s", err)
}
if !strings.Contains(string(data), strconv.Itoa(pid)) {
return errors.New("cannot cleanup process. process is not a part of this cgroup")
}
if err := ioutil.WriteFile(filepath.Join(basePath, procs), []byte(strconv.Itoa(pid)), 0644); err != nil {
return fmt.Errorf("cannot clean up process: %s", err)
}
return nil
}
// DeleteCgroup assumes the cgroup is already empty and destroys the directory structure.
// It will return an error if the group is not empty. Use RempoveProcess to remove all processes
// Before we try deletion
func (s *netCls) DeleteCgroup(cgroupname string) error {
_, err := os.Stat(filepath.Join(basePath, s.TriremePath, cgroupname))
if os.IsNotExist(err) {
zap.L().Debug("Group already deleted", zap.Error(err))
return nil
}
err = os.Remove(filepath.Join(basePath, s.TriremePath, cgroupname))
if err != nil {
return fmt.Errorf("unable to delete cgroup %s: %s", cgroupname, err)
}
return nil
}
//Deletebasepath removes the base aporeto directory which comes as a separate event when we are not managing any processes
func (s *netCls) Deletebasepath(cgroupName string) bool {
if cgroupName == s.TriremePath {
if err := os.Remove(filepath.Join(basePath, cgroupName)); err != nil {
zap.L().Error("Error when removing Trireme Base Path", zap.Error(err))
}
return true
}
return false
}
// ListCgroupProcesses returns lists of processes in the cgroup
func (s *netCls) ListCgroupProcesses(cgroupname string) ([]string, error) {
_, err := os.Stat(filepath.Join(basePath, s.TriremePath, cgroupname))
if os.IsNotExist(err) {
return []string{}, fmt.Errorf("cgroup %s does not exist: %s", cgroupname, err)
}
data, err := ioutil.ReadFile(filepath.Join(basePath, s.TriremePath, cgroupname, "cgroup.procs"))
if err != nil {
return []string{}, fmt.Errorf("cannot read procs file: %s", err)
}
procs := []string{}
for _, line := range strings.Split(string(data), "\n") {
if len(line) > 0 {
procs = append(procs, string(line))
}
}
return procs, nil
}
// ListAllCgroups returns a list of the cgroups that are managed in the Trireme path
func (s *netCls) ListAllCgroups(path string) []string {
cgroups, err := ioutil.ReadDir(filepath.Join(basePath, s.TriremePath, path))
if err != nil {
return []string{}
}
names := make([]string, len(cgroups))
for i := 0; i < len(cgroups); i++ {
names[i] = cgroups[i].Name()
}
return names
}
func mountCgroupController() {
mounts, err := ioutil.ReadFile("/proc/mounts")
if err != nil {
zap.L().Fatal(err.Error())
}
sc := bufio.NewScanner(strings.NewReader(string(mounts)))
var netCls = false
var cgroupMount string
for sc.Scan() {
if strings.HasPrefix(sc.Text(), "cgroup") {
cgroupMount = strings.Split(sc.Text(), " ")[1]
cgroupMount = cgroupMount[:strings.LastIndex(cgroupMount, "/")]
if strings.Contains(sc.Text(), "net_cls") {
basePath = strings.Split(sc.Text(), " ")[1]
netCls = true
return
}
}
}
if len(cgroupMount) == 0 {
zap.L().Error("Cgroups are not enabled or net_cls is not mounted")
return
}
if !netCls {
basePath = cgroupMount + "/net_cls"
if err := os.MkdirAll(basePath, 0700); err != nil {
zap.L().Fatal(err.Error())
}
if err := syscall.Mount("cgroup", basePath, "cgroup", 0, "net_cls,net_prio"); err != nil {
zap.L().Fatal(err.Error())
}
return
}
}
// CgroupMemberCount -- Returns the cound of the number of processes in a cgroup
func CgroupMemberCount(cgroupName string) int {
_, err := os.Stat(filepath.Join(basePath, TriremeBasePath, cgroupName))
if os.IsNotExist(err) {
return 0
}
data, err := ioutil.ReadFile(filepath.Join(basePath, TriremeBasePath, cgroupName, "cgroup.procs"))
if err != nil {
return 0
}
return len(data)
}
// NewDockerCgroupNetController returns a handle to call functions on the cgroup net_cls controller
func NewDockerCgroupNetController() Cgroupnetcls {
controller := &netCls{
markchan: make(chan uint64),
ReleaseAgentPath: "",
TriremePath: "",
}
return controller
}
//NewCgroupNetController returns a handle to call functions on the cgroup net_cls controller
func NewCgroupNetController(triremepath string, releasePath string) Cgroupnetcls {
binpath, _ := osext.Executable()
controller := &netCls{
markchan: make(chan uint64),
ReleaseAgentPath: binpath,
TriremePath: "",
}
if releasePath != "" {
controller.ReleaseAgentPath = releasePath
}
if triremepath != "" {
controller.TriremePath = triremepath
}
return controller
}
// MarkVal returns a new Mark Value
func MarkVal() uint64 {
return atomic.AddUint64(&markval, 1)
}