1 /* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
2 
3 #include "proxy.h"
4 
5 // mcp.add_stat(index, name)
6 // creates a custom lua stats counter
mcplib_add_stat(lua_State * L)7 int mcplib_add_stat(lua_State *L) {
8     int idx = luaL_checkinteger(L, -2);
9     const char *name = luaL_checkstring(L, -1);
10     proxy_ctx_t *ctx = PROXY_GET_CTX(L);
11 
12     if (idx < 1) {
13         proxy_lua_error(L, "stat index must be 1 or higher");
14         return 0;
15     }
16     if (idx > ctx->tunables.max_ustats) {
17         proxy_lua_ferror(L, "stat index must be %d or less", ctx->tunables.max_ustats);
18         return 0;
19     }
20     // max name length? avoids errors if something huge gets thrown in.
21     if (strlen(name) > STAT_KEY_LEN - 6) {
22         // we prepend "user_" to the output. + null byte.
23         proxy_lua_ferror(L, "stat name too long: %s\n", name);
24         return 0;
25     }
26     // restrict characters, at least no spaces/newlines.
27     for (int x = 0; x < strlen(name); x++) {
28         if (isspace(name[x])) {
29             proxy_lua_error(L, "stat cannot contain spaces or newlines");
30             return 0;
31         }
32     }
33 
34     STAT_L(ctx);
35     int stats_num = ctx->user_stats_num;
36     struct proxy_user_stats_entry *entries = ctx->user_stats;
37 
38     // if num_stats is 0 we need to init sizes.
39     // TODO (v2): malloc fail checking. (should be rare/impossible)
40     if (stats_num < idx) {
41         struct proxy_user_stats_entry *nentries = calloc(idx, sizeof(*entries));
42         // funny realloc; start with zeroed memory and copy in original.
43         if (entries) {
44             memcpy(nentries, entries, sizeof(*entries) * stats_num);
45             free(entries);
46         }
47         ctx->user_stats = nentries;
48         ctx->user_stats_num = idx;
49         entries = nentries;
50     }
51 
52     idx--; // real slot start as 0.
53     if (entries[idx].name != NULL) {
54         // If name changed, we have to reset the counter in the slot.
55         // Also only free/strdup the string if it's changed.
56         if (strcmp(name, entries[idx].name) != 0) {
57             entries[idx].reset = true;
58             free(entries[idx].name);
59             entries[idx].name = strdup(name);
60         }
61         // else the stat name didn't change, so don't do anything.
62     } else if (entries[idx].cname) {
63         char *oldname = ctx->user_stats_namebuf + entries[idx].cname;
64         if (strcmp(name, oldname) != 0) {
65             entries[idx].reset = true;
66             entries[idx].name = strdup(name);
67         }
68     } else {
69         entries[idx].name = strdup(name);
70     }
71     STAT_UL(ctx);
72 
73     return 0;
74 }
75 
mcplib_stat(lua_State * L)76 int mcplib_stat(lua_State *L) {
77     LIBEVENT_THREAD *t = PROXY_GET_THR(L);
78     if (t == NULL) {
79         proxy_lua_error(L, "stat must be called from router handlers");
80         return 0;
81     }
82 
83     struct proxy_user_stats *tus = t->proxy_user_stats;
84     if (tus == NULL) {
85         proxy_lua_error(L, "no stats counters initialized");
86         return 0;
87     }
88 
89     int idx = luaL_checkinteger(L, -2);
90     int change = luaL_checkinteger(L, -1);
91 
92     if (idx < 1 || idx > tus->num_stats) {
93         proxy_lua_error(L, "stat index out of range");
94         return 0;
95     }
96 
97     idx--; // actual array is 0 indexed.
98     WSTAT_L(t);
99     tus->counters[idx] += change;
100     WSTAT_UL(t);
101 
102     return 0;
103 }
104