• R/O
  • SSH
  • HTTPS

akari: 提交


Commit MetaInfo

修订版682 (tree)
时间2022-02-27 15:00:49
作者kumaneko

Log Message

Add forwardquota utility.

更改概述

差异

--- branches/forwardquota/Makefile (nonexistent)
+++ branches/forwardquota/Makefile (revision 682)
@@ -0,0 +1,19 @@
1+INSTALL := install
2+USRBINDIR := /usr/bin
3+ifndef CFLAGS
4+CFLAGS := -Wall -O2 -g -D_FORTIFY_SOURCE=2
5+endif
6+
7+all: forwardquota
8+
9+forwardquota: forwardquota.c
10+ $(CC) $(CPPFLAGS) $(CFLAGS) $(LDFLAGS) -o forwardquota forwardquota.c
11+
12+install: all
13+ mkdir -p -m 0755 $(INSTALLDIR)$(USRBINDIR)
14+ $(INSTALL) -m 0755 forwardquota $(INSTALLDIR)$(USRBINDIR)
15+
16+clean:
17+ rm -f -- forwardquota
18+
19+.PHONY: clean install
--- branches/forwardquota/forwardquota.c (nonexistent)
+++ branches/forwardquota/forwardquota.c (revision 682)
@@ -0,0 +1,134 @@
1+#include <stdio.h>
2+#include <string.h>
3+#include <sys/types.h>
4+#include <sys/stat.h>
5+#include <fcntl.h>
6+#include <unistd.h>
7+#include <sys/socket.h>
8+#include <arpa/inet.h>
9+#include <stdlib.h>
10+#include <sys/file.h>
11+#include <signal.h>
12+#include <sys/prctl.h>
13+#include <sys/poll.h>
14+#include <errno.h>
15+
16+static void forward_loop(const int client_fd, const int server_fd)
17+{
18+ const int fd[2] = { client_fd, server_fd };
19+ static char buffer[2][65536];
20+ int buffer_len[2] = { };
21+ _Bool read_eof[2] = { };
22+ int i;
23+ int len;
24+ while (1) {
25+ struct pollfd pfd[2] = {
26+ { fd[0], (buffer_len[0] ? POLLOUT : 0) | (buffer_len[1] < sizeof(buffer[1]) && !read_eof[0] ? POLLIN : 0) },
27+ { fd[1], (buffer_len[1] ? POLLOUT : 0) | (buffer_len[0] < sizeof(buffer[0]) && !read_eof[1] ? POLLIN : 0) }
28+ };
29+ if (pfd[0].events && pfd[1].events)
30+ poll(pfd, 2, -1);
31+ else if (pfd[0].events)
32+ poll(pfd, 1, -1);
33+ else if (pfd[1].events)
34+ poll(&pfd[1], 1, -1);
35+ else
36+ break;
37+ for (i = 0; i < 2; i++) {
38+ if (!(pfd[i].revents & POLLIN))
39+ continue;
40+ len = recv(fd[i], buffer[i ^ 1] + buffer_len[i ^ 1], sizeof(buffer[0]) - buffer_len[i ^ 1], MSG_DONTWAIT);
41+ if (len > 0) {
42+ buffer_len[i ^ 1] += len;
43+ pfd[i ^ 1].revents |= POLLOUT;
44+ } else if (len == 0) {
45+ read_eof[i] = 1;
46+ if (!buffer_len[i ^ 1])
47+ shutdown(fd[i ^ 1], SHUT_WR);
48+ }
49+ }
50+ for (i = 0; i < 2; i++) {
51+ if (!(pfd[i].revents & POLLOUT))
52+ continue;
53+ len = send(fd[i], buffer[i], buffer_len[i], MSG_DONTWAIT);
54+ if (len > 0) {
55+ buffer_len[i] -= len;
56+ if (buffer_len[i])
57+ memmove(buffer[i], buffer[i] + len, buffer_len[i]);
58+ else if (read_eof[i ^ 1])
59+ shutdown(fd[i], SHUT_WR);
60+ }
61+ }
62+ }
63+}
64+
65+int main(int argc, char *argv[])
66+{
67+ struct sockaddr_in addr1 = { .sin_family = AF_INET, .sin_addr.s_addr = htonl(INADDR_ANY) };
68+ struct sockaddr_in addr2 = { .sin_family = AF_INET, .sin_addr.s_addr = htonl(INADDR_LOOPBACK) };
69+ const int listener_fd = socket(PF_INET, SOCK_STREAM, 0);
70+ if (argc != 4)
71+ goto usage;
72+ addr1.sin_port = htons(atoi(argv[1]));
73+ if (!addr1.sin_port || bind(listener_fd, (struct sockaddr *) &addr1, sizeof(addr1)) || listen(listener_fd, 5)) {
74+ fprintf(stderr, "Can't listen on 0.0.0.0:%u.\n", ntohs(addr1.sin_port));
75+ return 1;
76+ } else if (chdir(argv[2])) {
77+ fprintf(stderr, "Can't change directory to %s .\n", argv[2]);
78+ return 1;
79+ }
80+ addr2.sin_port = htons(atoi(argv[3]));
81+ if (!addr2.sin_port) {
82+ fprintf(stderr, "Can't connect to 127.0.0.1:%u.\n", ntohs(addr2.sin_port));
83+ return 1;
84+ }
85+ signal(SIGCLD, SIG_IGN);
86+ while (1) {
87+ socklen_t size = sizeof(addr1);
88+ const int client_fd = accept(listener_fd, (struct sockaddr *) &addr1, &size);
89+ if (client_fd == EOF) {
90+ sleep(1);
91+ continue;
92+ }
93+ if (fork() == 0) {
94+ static char buffer[128] = { };
95+ static char ipbuf[40] = { };
96+ int idx = 0;
97+ int lock_fd = EOF;
98+ const int server_fd = socket(PF_INET, SOCK_STREAM, 0);
99+ close(listener_fd);
100+ inet_ntop(AF_INET, &addr1.sin_addr, ipbuf, sizeof(ipbuf) - 1);
101+ prctl(PR_SET_NAME, (unsigned long) ipbuf, 0, 0, 0);
102+ while (1) {
103+ snprintf(buffer, sizeof(buffer) - 1, "%s-%u", ipbuf, idx++);
104+ lock_fd = open(buffer, O_RDONLY);
105+ if (lock_fd == EOF && errno == ENOENT)
106+ _exit(1);
107+ if (flock(lock_fd, LOCK_EX | LOCK_NB) == 0)
108+ break;
109+ close(lock_fd);
110+ }
111+ if (connect(server_fd, (struct sockaddr *) &addr2, sizeof(addr2)))
112+ _exit(1);
113+ forward_loop(client_fd, server_fd);
114+ _exit(0);
115+ }
116+ if (close(client_fd)) {
117+ fprintf(stderr, "Can't close socket.\n");
118+ return 1;
119+ }
120+ }
121+usage:
122+ fprintf(stderr, "Simple quota for TCP/IPv4 connections.\n\n"
123+ "Usage:\n %s $listen_port $quota_dir $forward_port\n\n", argv[0]);
124+ fputs("This program allows restricting number of connections from the same client.\n\n"
125+ "This program listens on 0.0.0.0:$listen_port and accepts connections from any address, "
126+ "but forwards that connection to 127.0.0.1:$forward_port only if successfully grabbed "
127+ "one of lock files for that client in the $quota_dir directory.\n\n"
128+ "The lock files must be readable and have $ip-$index filename where $ip is the client's "
129+ "IPv4 address and $index is a sequential integer starting from 0. "
130+ "For example, if only 192.168.1.100-0, 192.168.1.100-1 and 192.168.1.101-0 are there, "
131+ "192.168.1.100 can forward up to 2 connections, 192.168.1.101 can forward up to 1 "
132+ "connection, and all other clients can't forward connection.\n", stderr);
133+ return 1;
134+}
Show on old repository browser