Hugging Face's logo

Spaces:
Moge-Row
/
Row-proxy
Paused

App Files Community
4
Moge-Row commited on
Commit
bf9eaec
·
1 Parent(s): 22c8f3c

fix key lookup bug

Browse files
Files changed (1) hide show
  1. internal/handler/keys.go +37 -42
internal/handler/keys.go CHANGED
@@ -21,6 +21,7 @@ type Store struct {
21
  }
22
 
23
  var s = &Store{Keys: make(map[string]*KeyData)}
 
24
  const p = "/data/keys.json"
25
 
26
  func init() {
@@ -35,7 +36,9 @@ func init() {
35
  save()
36
  }
37
  }
38
- if s.Keys == nil { s.Keys = make(map[string]*KeyData) }
 
 
39
  }
40
 
41
  func save() {
@@ -44,24 +47,33 @@ func save() {
44
  }
45
 
46
  func TrackUsage(k string, t int) {
47
- if k == "free" { return }
 
 
48
  s.mu.Lock()
49
  defer s.mu.Unlock()
50
- if kd, ok := s.Keys[k]; ok {
51
- kd.Tokens += t
52
- save()
 
 
 
53
  }
54
  }
55
 
56
  func CheckAndTrack(k string, t int) (bool, string) {
57
- if k == "free" { return true, "" }
 
 
58
  s.mu.Lock()
59
  defer s.mu.Unlock()
60
- if kd, ok := s.Keys[k]; ok {
61
- kd.Requests++
62
- kd.Tokens += t
63
- save()
64
- return true, ""
 
 
65
  }
66
  return false, "Key invalida"
67
  }
@@ -69,27 +81,27 @@ func CheckAndTrack(k string, t int) (bool, string) {
69
  func GenerateKey(n string) string {
70
  const chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
71
  b := make([]byte, 8)
72
- for i := range b { b[i] = chars[rand.Intn(len(chars))] }
 
 
73
  k := "RWPX-" + string(b)
74
-
75
  s.mu.Lock()
76
  s.Keys[n+":"+k] = &KeyData{Name: n, Key: k}
77
  s.mu.Unlock()
78
-
79
  go func(name, key string) {
80
  BackupKeyToSheet(name, key)
81
  }(n, k)
82
-
83
  save()
84
  return k
85
  }
86
 
87
  func HandleGenKey(w http.ResponseWriter, r *http.Request) {
88
  n := r.FormValue("name")
89
- if n == "" { n = "User" }
 
 
90
  existingKey := r.FormValue("key")
91
  var k string
92
-
93
  if existingKey != "" {
94
  s.mu.Lock()
95
  s.Keys[n+":"+existingKey] = &KeyData{Name: n, Key: existingKey}
@@ -102,7 +114,6 @@ func HandleGenKey(w http.ResponseWriter, r *http.Request) {
102
  } else {
103
  k = GenerateKey(n)
104
  }
105
-
106
  w.Header().Set("Content-Type", "application/json")
107
  resp := map[string]string{"key": k, "name": n}
108
  json.NewEncoder(w).Encode(resp)
@@ -120,33 +131,17 @@ func HandleDeleteKey(w http.ResponseWriter, r *http.Request) {
120
  }
121
  s.mu.Lock()
122
  defer s.mu.Unlock()
123
- delete(s.Keys, k)
124
- save()
 
 
 
 
 
125
  w.Write([]byte("OK"))
126
  }
127
 
128
  func HandleStats(w http.ResponseWriter, r *http.Request) {
129
  s.mu.Lock()
130
  defer s.mu.Unlock()
131
- w.Header().Set("Content-Type", "application/json")
132
- public := make(map[string]interface{})
133
- for _, v := range s.Keys {
134
- public[v.Name] = map[string]interface{}{"name": v.Name, "requests": v.Requests}
135
- }
136
- json.NewEncoder(w).Encode(map[string]interface{}{"keys": public})
137
- }
138
-
139
- func HandleSecretReveal(w http.ResponseWriter, r *http.Request) {
140
- if r.URL.Query().Get("token") != "CCCP_IS2_1944_STALIN" {
141
- w.WriteHeader(403)
142
- return
143
- }
144
- s.mu.Lock()
145
- defer s.mu.Unlock()
146
- w.Header().Set("Content-Type", "text/html")
147
- w.Write([]byte("<body style='background:#000;color:#0f0;font-family:monospace;'>"))
148
- for compositeKey := range s.Keys {
149
- w.Write([]byte("<p>" + compositeKey + "</p>"))
150
- }
151
- w.Write([]byte("</body>"))
152
- }
 
21
  }
22
 
23
  var s = &Store{Keys: make(map[string]*KeyData)}
24
+
25
  const p = "/data/keys.json"
26
 
27
  func init() {
 
36
  save()
37
  }
38
  }
39
+ if s.Keys == nil {
40
+ s.Keys = make(map[string]*KeyData)
41
+ }
42
  }
43
 
44
  func save() {
 
47
  }
48
 
49
  func TrackUsage(k string, t int) {
50
+ if k == "free" {
51
+ return
52
+ }
53
  s.mu.Lock()
54
  defer s.mu.Unlock()
55
+ for _, kd := range s.Keys {
56
+ if kd.Key == k {
57
+ kd.Tokens += t
58
+ save()
59
+ return
60
+ }
61
  }
62
  }
63
 
64
  func CheckAndTrack(k string, t int) (bool, string) {
65
+ if k == "free" {
66
+ return true, ""
67
+ }
68
  s.mu.Lock()
69
  defer s.mu.Unlock()
70
+ for _, kd := range s.Keys {
71
+ if kd.Key == k {
72
+ kd.Requests++
73
+ kd.Tokens += t
74
+ save()
75
+ return true, ""
76
+ }
77
  }
78
  return false, "Key invalida"
79
  }
 
81
  func GenerateKey(n string) string {
82
  const chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
83
  b := make([]byte, 8)
84
+ for i := range b {
85
+ b[i] = chars[rand.Intn(len(chars))]
86
+ }
87
  k := "RWPX-" + string(b)
 
88
  s.mu.Lock()
89
  s.Keys[n+":"+k] = &KeyData{Name: n, Key: k}
90
  s.mu.Unlock()
 
91
  go func(name, key string) {
92
  BackupKeyToSheet(name, key)
93
  }(n, k)
 
94
  save()
95
  return k
96
  }
97
 
98
  func HandleGenKey(w http.ResponseWriter, r *http.Request) {
99
  n := r.FormValue("name")
100
+ if n == "" {
101
+ n = "User"
102
+ }
103
  existingKey := r.FormValue("key")
104
  var k string
 
105
  if existingKey != "" {
106
  s.mu.Lock()
107
  s.Keys[n+":"+existingKey] = &KeyData{Name: n, Key: existingKey}
 
114
  } else {
115
  k = GenerateKey(n)
116
  }
 
117
  w.Header().Set("Content-Type", "application/json")
118
  resp := map[string]string{"key": k, "name": n}
119
  json.NewEncoder(w).Encode(resp)
 
131
  }
132
  s.mu.Lock()
133
  defer s.mu.Unlock()
134
+ for compositeKey, kd := range s.Keys {
135
+ if kd.Key == k {
136
+ delete(s.Keys, compositeKey)
137
+ save()
138
+ break
139
+ }
140
+ }
141
  w.Write([]byte("OK"))
142
  }
143
 
144
  func HandleStats(w http.ResponseWriter, r *http.Request) {
145
  s.mu.Lock()
146
  defer s.mu.Unlock()
147
+ w.Header(