This repository has been archived by the owner on Feb 12, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
/
user.go
108 lines (87 loc) · 2.71 KB
/
user.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
package routes
import (
"encoding/json"
"fmt"
"github.com/emicklei/go-restful/v3"
"github.com/ractf/andromeda/pkg/node"
"github.com/ractf/andromeda/pkg/node/instance"
"math/rand"
"net/http"
"sync"
)
type userRoutes struct {
node *node.Node
users map[string]*instance.Instance
mutex *sync.Mutex
}
func AddUserRoutes(node *node.Node, ws *restful.WebService) {
u := userRoutes{
node: node,
users: make(map[string]*instance.Instance),
mutex: &sync.Mutex{},
}
ws.Route(ws.POST("/").To(Authenticated(u.getInstance, node.Config.ApiKey)))
ws.Route(ws.POST("/reset").To(Authenticated(u.resetInstance, node.Config.ApiKey)))
}
type instanceRequest struct {
Job string `json:"job"`
User string `json:"user"`
}
type instanceDetails struct {
Ip string `json:"ip"`
Port int `json:"port"`
}
func (u *userRoutes) getInstance(request *restful.Request, response *restful.Response) {
request.Request.Body = http.MaxBytesReader(response.ResponseWriter, request.Request.Body, 10240)
dec := json.NewDecoder(request.Request.Body)
var r instanceRequest
err := dec.Decode(&r)
if err != nil {
fmt.Println(err)
response.WriteError(500, err)
return
}
job := r.Job
jobSpec := u.node.GetJobSpecByName(job)
u.mutex.Lock()
defer u.mutex.Unlock()
if instance, ok := u.users[r.User]; ok {
if instance.Job == jobSpec && !instance.Dead {
_ = response.WriteAsJson(instanceDetails{Port: instance.Port, Ip: u.node.Config.PublicIp})
return
}
}
instances := u.node.InstanceController.GetLocalInstancesOf(jobSpec)
rand.Shuffle(len(instances), func(i, j int) { instances[i], instances[j] = instances[j], instances[i] })
u.users[r.User] = instances[0]
_ = response.WriteAsJson(instanceDetails{Port: instances[0].Port, Ip: u.node.Config.PublicIp})
}
func (u *userRoutes) resetInstance(request *restful.Request, response *restful.Response) {
request.Request.Body = http.MaxBytesReader(response.ResponseWriter, request.Request.Body, 10240)
dec := json.NewDecoder(request.Request.Body)
var r instanceRequest
err := dec.Decode(&r)
if err != nil {
fmt.Println(err)
response.WriteError(500, err)
return
}
job := r.Job
jobSpec := u.node.GetJobSpecByName(job)
u.mutex.Lock()
defer u.mutex.Unlock()
var avoid *instance.Instance = nil
if instance, ok := u.users[r.User]; ok {
if instance.Job == jobSpec {
avoid = instance
}
}
instances := u.node.InstanceController.GetLocalInstancesOf(jobSpec)
rand.Shuffle(len(instances), func(i, j int) { instances[i], instances[j] = instances[j], instances[i] })
index := 0
if instances[0] == avoid && len(instances) > 1 {
index = 1
}
u.users[r.User] = instances[index]
_ = response.WriteAsJson(instanceDetails{Port: instances[index].Port, Ip: u.node.Config.PublicIp})
}