805 lines
26 KiB
C
805 lines
26 KiB
C
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <unistd.h>
|
|
#include <dlfcn.h>
|
|
#include <sys/types.h>
|
|
#include <sys/stat.h>
|
|
#include <fcntl.h>
|
|
#include <time.h>
|
|
#include <sys/uio.h>
|
|
#include <sys/sendfile.h>
|
|
#include <linux/fs.h>
|
|
#include <string.h>
|
|
#include <stdarg.h>
|
|
#include <sys/socket.h>
|
|
|
|
// Function pointers for original read and write and other relevant calls
|
|
|
|
// Function pointers for standard I/O functions
|
|
int (*original_putchar_unlocked)(int);
|
|
int (*original_putc)(int, FILE *);
|
|
int (*original_putc_unlocked)(int, FILE *);
|
|
size_t (*original_fread)(void *, size_t, size_t, FILE *);
|
|
size_t (*original_fread_unlocked)(void *, size_t, size_t, FILE *);
|
|
int (*original_fgetc)(FILE *);
|
|
char *(*original_fgets)(char *, int, FILE *);
|
|
int (*original_getc)(FILE *);
|
|
int (*original_getc_unlocked)(FILE *);
|
|
|
|
// Function pointers for socket functions
|
|
ssize_t (*original_recv)(int, void *, size_t, int);
|
|
ssize_t (*original_recvfrom)(int, void *, size_t, int, struct sockaddr *, socklen_t *);
|
|
ssize_t (*original_recvmsg)(int, struct msghdr *, int);
|
|
|
|
// Function pointers for formatted input functions
|
|
int (*original_scanf)(const char *, ...);
|
|
int (*original_fscanf)(FILE *, const char *, ...);
|
|
|
|
// Additional low-level function pointers
|
|
ssize_t (*original_pread)(int fd, void *buf, size_t count, off_t offset);
|
|
ssize_t (*original_pwrite)(int fd, const void *buf, size_t count, off_t offset);
|
|
ssize_t (*original_readv)(int fd, const struct iovec *iov, int iovcnt);
|
|
ssize_t (*original_writev)(int fd, const struct iovec *iov, int iovcnt);
|
|
|
|
|
|
ssize_t (*original_read)(int fd, void *buf, size_t count);
|
|
ssize_t (*original_write)(int fd, const void *buf, size_t count);
|
|
size_t (*original_fwrite)(const void *ptr, size_t size, size_t nmemb, FILE *stream);
|
|
size_t (*original_fwrite_unlocked)(const void *ptr, size_t size, size_t nmemb, FILE *stream);
|
|
int (*original_fflush)(FILE *stream);
|
|
int (*original_fputs)(const char *str, FILE *stream);
|
|
int (*original_fputs_unlocked)(const char *str, FILE *stream);
|
|
int (*original_putchar)(int c);
|
|
int (*original_dprintf)(int fd, const char *format, ...);
|
|
ssize_t (*original_pread)(int fd, void *buf, size_t count, off_t offset);
|
|
ssize_t (*original_pwrite)(int fd, const void *buf, size_t count, off_t offset);
|
|
ssize_t (*original_readv)(int fd, const struct iovec *iov, int iovcnt);
|
|
ssize_t (*original_writev)(int fd, const struct iovec *iov, int iovcnt);
|
|
off_t (*original_sendfile)(int out_fd, int in_fd, off_t *offset, size_t count);
|
|
ssize_t (*original_splice)(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags);
|
|
ssize_t (*original_copy_file_range)(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags);
|
|
|
|
// Function pointer for original functions
|
|
ssize_t (*original_send)(int sockfd, const void *buf, size_t len, int flags);
|
|
ssize_t (*original_sendto)(int sockfd, const void *buf, size_t len, int flags,
|
|
const struct sockaddr *dest_addr, socklen_t addrlen);
|
|
ssize_t (*original_sendmsg)(int sockfd, const struct msghdr *msg, int flags);
|
|
ssize_t (*original_vmsplice)(int fd, const struct iovec *iov, unsigned long nr_segs,
|
|
unsigned int flags);
|
|
ssize_t (*original_process_vm_readv)(pid_t pid, const struct iovec *local_iov,
|
|
int local_iovcnt, const struct iovec *remote_iov,
|
|
int remote_iovcnt, unsigned long flags);
|
|
ssize_t (*original_process_vm_writev)(pid_t pid, const struct iovec *local_iov,
|
|
int local_iovcnt, const struct iovec *remote_iov,
|
|
int remote_iovcnt, unsigned long flags);
|
|
|
|
// Rate limits in bytes per second
|
|
static size_t read_limit = 0; // 0 = no limit
|
|
static size_t write_limit = 0;
|
|
static int debug_mode = 0;
|
|
|
|
// Tracking variables
|
|
static size_t read_bytes = 0;
|
|
static size_t write_bytes = 0;
|
|
static struct timespec read_start, write_start;
|
|
|
|
void throttle(size_t *bytes, size_t limit, struct timespec *start) {
|
|
if (limit == 0) {
|
|
return; // No limit applied
|
|
}
|
|
|
|
struct timespec now;
|
|
clock_gettime(CLOCK_MONOTONIC, &now);
|
|
|
|
double elapsed = (now.tv_sec - start->tv_sec) + (now.tv_nsec - start->tv_nsec) / 1e9;
|
|
// If it is > 1 second since last write:
|
|
// treat as if it is only 1 second
|
|
elapsed = elapsed > 1 ? 1 : elapsed;
|
|
double allowed_bytes = limit * elapsed;
|
|
if (debug_mode) fprintf(stderr, "allowed: %ld written: %ld\n",
|
|
(long)allowed_bytes, *bytes);
|
|
if (*bytes > allowed_bytes) {
|
|
*bytes -= allowed_bytes;
|
|
double sleep_time = *bytes / (double)limit;
|
|
|
|
if (debug_mode) fprintf(stderr, "sleep: %f\n", sleep_time);
|
|
usleep((useconds_t)(sleep_time * 1e6));
|
|
*start = now;
|
|
}
|
|
}
|
|
|
|
size_t fwrite(const void *ptr, size_t size, size_t nmemb, FILE *stream) {
|
|
if (!original_fwrite) {
|
|
original_fwrite = dlsym(RTLD_NEXT, "fwrite");
|
|
if (!original_fwrite) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original fwrite: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = size * nmemb;
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
size_t result = original_fwrite(ptr, size, nmemb, stream);
|
|
if (result > 0) {
|
|
write_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
#undef fwrite_unlocked
|
|
size_t fwrite_unlocked(const void *ptr, size_t size, size_t nmemb, FILE *stream) {
|
|
if (!original_fwrite_unlocked) {
|
|
original_fwrite_unlocked = dlsym(RTLD_NEXT, "fwrite_unlocked");
|
|
if (!original_fwrite_unlocked) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original fwrite_unlocked: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = size * nmemb;
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
size_t result = original_fwrite_unlocked(ptr, size, nmemb, stream);
|
|
if (result > 0) {
|
|
write_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int fputs(const char *str, FILE *stream) {
|
|
if (!original_fputs) {
|
|
original_fputs = dlsym(RTLD_NEXT, "fputs");
|
|
if (!original_fputs) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original fputs: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = strlen(str);
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
int result = original_fputs(str, stream);
|
|
if (result >= 0) {
|
|
write_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int fputs_unlocked(const char *str, FILE *stream) {
|
|
if (!original_fputs_unlocked) {
|
|
original_fputs_unlocked = dlsym(RTLD_NEXT, "fputs_unlocked");
|
|
if (!original_fputs_unlocked) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original fputs_unlocked: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = strlen(str);
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
int result = original_fputs_unlocked(str, stream);
|
|
if (result >= 0) {
|
|
write_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int putchar(int c) {
|
|
if (!original_putchar) {
|
|
original_putchar = dlsym(RTLD_NEXT, "putchar");
|
|
if (!original_putchar) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original putchar: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
int result = original_putchar(c);
|
|
if (result != EOF) {
|
|
write_bytes += 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int dprintf(int fd, const char *format, ...) {
|
|
if (!original_dprintf) {
|
|
original_dprintf = dlsym(RTLD_NEXT, "dprintf");
|
|
if (!original_dprintf) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original dprintf: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
va_list args;
|
|
va_start(args, format);
|
|
char buffer[1024];
|
|
int len = vsnprintf(buffer, sizeof(buffer), format, args);
|
|
va_end(args);
|
|
|
|
if (len > 0) {
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
write_bytes += len;
|
|
}
|
|
|
|
return original_dprintf(fd, "%s", buffer);
|
|
}
|
|
|
|
int fflush(FILE *stream) {
|
|
if (!original_fflush) {
|
|
original_fflush = dlsym(RTLD_NEXT, "fflush");
|
|
if (!original_fflush) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original fflush: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
return original_fflush(stream);
|
|
}
|
|
|
|
ssize_t read(int fd, void *buf, size_t count) {
|
|
if (!original_read) {
|
|
original_read = dlsym(RTLD_NEXT, "read");
|
|
if (!original_read) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original read: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
|
|
ssize_t result = original_read(fd, buf, count);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t write(int fd, const void *buf, size_t count) {
|
|
if (!original_write) {
|
|
original_write = dlsym(RTLD_NEXT, "write");
|
|
if (!original_write) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original write: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
ssize_t result = original_write(fd, buf, count);
|
|
if (result > 0) {
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t pread(int fd, void *buf, size_t count, off_t offset) {
|
|
if (!original_pread) {
|
|
original_pread = dlsym(RTLD_NEXT, "pread");
|
|
if (!original_pread) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original pread: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
|
|
ssize_t result = original_pread(fd, buf, count, offset);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t pwrite(int fd, const void *buf, size_t count, off_t offset) {
|
|
if (!original_pwrite) {
|
|
original_pwrite = dlsym(RTLD_NEXT, "pwrite");
|
|
if (!original_pwrite) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original pwrite: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
ssize_t result = original_pwrite(fd, buf, count, offset);
|
|
if (result > 0) {
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t readv(int fd, const struct iovec *iov, int iovcnt) {
|
|
if (!original_readv) {
|
|
original_readv = dlsym(RTLD_NEXT, "readv");
|
|
if (!original_readv) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original readv: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
|
|
ssize_t result = original_readv(fd, iov, iovcnt);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t writev(int fd, const struct iovec *iov, int iovcnt) {
|
|
if (!original_writev) {
|
|
original_writev = dlsym(RTLD_NEXT, "writev");
|
|
if (!original_writev) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original writev: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
ssize_t result = original_writev(fd, iov, iovcnt);
|
|
if (result > 0) {
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
off_t sendfile(int out_fd, int in_fd, off_t *offset, size_t count) {
|
|
if (!original_sendfile) {
|
|
original_sendfile = dlsym(RTLD_NEXT, "sendfile");
|
|
if (!original_sendfile) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original sendfile: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
off_t result = original_sendfile(out_fd, in_fd, offset, count);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t splice(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags) {
|
|
if (!original_splice) {
|
|
original_splice = dlsym(RTLD_NEXT, "splice");
|
|
if (!original_splice) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original splice: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
ssize_t result = original_splice(fd_in, off_in, fd_out, off_out, len, flags);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t copy_file_range(int fd_in, loff_t *off_in, int fd_out, loff_t *off_out, size_t len, unsigned int flags) {
|
|
if (!original_copy_file_range) {
|
|
original_copy_file_range = dlsym(RTLD_NEXT, "copy_file_range");
|
|
if (!original_copy_file_range) {
|
|
if (debug_mode) fprintf(stderr, "Error loading original copy_file_range: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
|
|
ssize_t result = original_copy_file_range(fd_in, off_in, fd_out, off_out, len, flags);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int putchar_unlocked(int c) {
|
|
if (!original_putchar_unlocked) {
|
|
original_putchar_unlocked = dlsym(RTLD_NEXT, "putchar_unlocked");
|
|
if (!original_putchar_unlocked) {
|
|
fprintf(stderr, "Error loading original putchar_unlocked: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
int result = original_putchar_unlocked(c);
|
|
if (result != EOF) {
|
|
write_bytes += 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t recv(int sockfd, void *buf, size_t len, int flags) {
|
|
if (!original_recv) {
|
|
original_recv = dlsym(RTLD_NEXT, "recv");
|
|
if (!original_recv) {
|
|
fprintf(stderr, "Error loading original recv: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
ssize_t result = original_recv(sockfd, buf, len, flags);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t recvfrom(int sockfd, void *buf, size_t len, int flags,
|
|
struct sockaddr *src_addr, socklen_t *addrlen) {
|
|
if (!original_recvfrom) {
|
|
original_recvfrom = dlsym(RTLD_NEXT, "recvfrom");
|
|
if (!original_recvfrom) {
|
|
fprintf(stderr, "Error loading original recvfrom: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
ssize_t result = original_recvfrom(sockfd, buf, len, flags, src_addr, addrlen);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t recvmsg(int sockfd, struct msghdr *msg, int flags) {
|
|
if (!original_recvmsg) {
|
|
original_recvmsg = dlsym(RTLD_NEXT, "recvmsg");
|
|
if (!original_recvmsg) {
|
|
fprintf(stderr, "Error loading original recvmsg: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
ssize_t result = original_recvmsg(sockfd, msg, flags);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int putc(int c, FILE *stream) {
|
|
if (!original_putc) {
|
|
original_putc = dlsym(RTLD_NEXT, "putc");
|
|
if (!original_putc) {
|
|
fprintf(stderr, "Error loading original putc: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
int result = original_putc(c, stream);
|
|
if (result != EOF) {
|
|
write_bytes += 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int putc_unlocked(int c, FILE *stream) {
|
|
if (!original_putc_unlocked) {
|
|
original_putc_unlocked = dlsym(RTLD_NEXT, "putc_unlocked");
|
|
if (!original_putc_unlocked) {
|
|
fprintf(stderr, "Error loading original putc_unlocked: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
int result = original_putc_unlocked(c, stream);
|
|
if (result != EOF) {
|
|
write_bytes += 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
size_t fread(void *ptr, size_t size, size_t nmemb, FILE *stream) {
|
|
if (!original_fread) {
|
|
original_fread = dlsym(RTLD_NEXT, "fread");
|
|
if (!original_fread) {
|
|
fprintf(stderr, "Error loading original fread: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
size_t result = original_fread(ptr, size, nmemb, stream);
|
|
if (result > 0) {
|
|
read_bytes += result * size;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
size_t fread_unlocked(void *ptr, size_t size, size_t nmemb, FILE *stream) {
|
|
if (!original_fread_unlocked) {
|
|
original_fread_unlocked = dlsym(RTLD_NEXT, "fread_unlocked");
|
|
if (!original_fread_unlocked) {
|
|
fprintf(stderr, "Error loading original fread_unlocked: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
size_t result = original_fread_unlocked(ptr, size, nmemb, stream);
|
|
if (result > 0) {
|
|
read_bytes += result * size;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int fgetc(FILE *stream) {
|
|
if (!original_fgetc) {
|
|
original_fgetc = dlsym(RTLD_NEXT, "fgetc");
|
|
if (!original_fgetc) {
|
|
fprintf(stderr, "Error loading original fgetc: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
int result = original_fgetc(stream);
|
|
if (result != EOF) {
|
|
read_bytes += 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
char *fgets(char *s, int size, FILE *stream) {
|
|
if (!original_fgets) {
|
|
original_fgets = dlsym(RTLD_NEXT, "fgets");
|
|
if (!original_fgets) {
|
|
fprintf(stderr, "Error loading original fgets: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
char *result = original_fgets(s, size, stream);
|
|
if (result != NULL) {
|
|
read_bytes += strlen(result);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int getc(FILE *stream) {
|
|
if (!original_getc) {
|
|
original_getc = dlsym(RTLD_NEXT, "getc");
|
|
if (!original_getc) {
|
|
fprintf(stderr, "Error loading original getc: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
int result = original_getc(stream);
|
|
if (result != EOF) {
|
|
read_bytes += 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int getc_unlocked(FILE *stream) {
|
|
if (!original_getc_unlocked) {
|
|
original_getc_unlocked = dlsym(RTLD_NEXT, "getc_unlocked");
|
|
if (!original_getc_unlocked) {
|
|
fprintf(stderr, "Error loading original getc_unlocked: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
int result = original_getc_unlocked(stream);
|
|
if (result != EOF) {
|
|
read_bytes += 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int scanf(const char *format, ...) {
|
|
if (!original_scanf) {
|
|
original_scanf = dlsym(RTLD_NEXT, "scanf");
|
|
if (!original_scanf) {
|
|
fprintf(stderr, "Error loading original scanf: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start); // Throttle before reading
|
|
va_list args;
|
|
va_start(args, format);
|
|
int result = vscanf(format, args);
|
|
va_end(args);
|
|
if (result > 0) {
|
|
// Note: Accurate byte count for scanf is difficult; this is an estimate
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
int fscanf(FILE *stream, const char *format, ...) {
|
|
if (!original_fscanf) {
|
|
original_fscanf = dlsym(RTLD_NEXT, "fscanf");
|
|
if (!original_fscanf) {
|
|
fprintf(stderr, "Error loading original fscanf: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
va_list args;
|
|
va_start(args, format);
|
|
int result = vfscanf(stream, format, args);
|
|
va_end(args);
|
|
if (result > 0) {
|
|
read_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t send(int sockfd, const void *buf, size_t len, int flags) {
|
|
if (!original_send) {
|
|
original_send = dlsym(RTLD_NEXT, "send");
|
|
if (!original_send) {
|
|
fprintf(stderr, "Error loading original send: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
ssize_t result = original_send(sockfd, buf, len, flags);
|
|
if (result > 0) {
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t sendto(int sockfd, const void *buf, size_t len, int flags,
|
|
const struct sockaddr *dest_addr, socklen_t addrlen) {
|
|
if (!original_sendto) {
|
|
original_sendto = dlsym(RTLD_NEXT, "sendto");
|
|
if (!original_sendto) {
|
|
fprintf(stderr, "Error loading original sendto: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
ssize_t result = original_sendto(sockfd, buf, len, flags, dest_addr, addrlen);
|
|
if (result > 0) {
|
|
write_bytes += result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t sendmsg(int sockfd, const struct msghdr *msg, int flags) {
|
|
if (!original_sendmsg) {
|
|
original_sendmsg = dlsym(RTLD_NEXT, "sendmsg");
|
|
if (!original_sendmsg) {
|
|
fprintf(stderr, "Error loading original sendmsg: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = 0;
|
|
for (int i = 0; i < msg->msg_iovlen; i++) {
|
|
total_bytes += msg->msg_iov[i].iov_len;
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
ssize_t result = original_sendmsg(sockfd, msg, flags);
|
|
if (result > 0) {
|
|
write_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t vmsplice(int fd, const struct iovec *iov, unsigned long nr_segs,
|
|
unsigned int flags) {
|
|
if (!original_vmsplice) {
|
|
original_vmsplice = dlsym(RTLD_NEXT, "vmsplice");
|
|
if (!original_vmsplice) {
|
|
fprintf(stderr, "Error loading original vmsplice: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = 0;
|
|
for (unsigned long i = 0; i < nr_segs; i++) {
|
|
total_bytes += iov[i].iov_len;
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
ssize_t result = original_vmsplice(fd, iov, nr_segs, flags);
|
|
if (result > 0) {
|
|
write_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t process_vm_readv(pid_t pid, const struct iovec *local_iov, int local_iovcnt,
|
|
const struct iovec *remote_iov, int remote_iovcnt,
|
|
unsigned long flags) {
|
|
if (!original_process_vm_readv) {
|
|
original_process_vm_readv = dlsym(RTLD_NEXT, "process_vm_readv");
|
|
if (!original_process_vm_readv) {
|
|
fprintf(stderr, "Error loading original process_vm_readv: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = 0;
|
|
for (int i = 0; i < local_iovcnt; i++) {
|
|
total_bytes += local_iov[i].iov_len;
|
|
}
|
|
|
|
throttle(&read_bytes, read_limit, &read_start);
|
|
ssize_t result = original_process_vm_readv(pid, local_iov, local_iovcnt,
|
|
remote_iov, remote_iovcnt, flags);
|
|
if (result > 0) {
|
|
read_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
ssize_t process_vm_writev(pid_t pid, const struct iovec *local_iov, int local_iovcnt,
|
|
const struct iovec *remote_iov, int remote_iovcnt,
|
|
unsigned long flags) {
|
|
if (!original_process_vm_writev) {
|
|
original_process_vm_writev = dlsym(RTLD_NEXT, "process_vm_writev");
|
|
if (!original_process_vm_writev) {
|
|
fprintf(stderr, "Error loading original process_vm_writev: %s\n", dlerror());
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
}
|
|
|
|
size_t total_bytes = 0;
|
|
for (int i = 0; i < local_iovcnt; i++) {
|
|
total_bytes += local_iov[i].iov_len;
|
|
}
|
|
|
|
throttle(&write_bytes, write_limit, &write_start);
|
|
ssize_t result = original_process_vm_writev(pid, local_iov, local_iovcnt,
|
|
remote_iov, remote_iovcnt, flags);
|
|
if (result > 0) {
|
|
write_bytes += total_bytes;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
__attribute__((constructor))
|
|
void init() {
|
|
clock_gettime(CLOCK_MONOTONIC, &read_start);
|
|
clock_gettime(CLOCK_MONOTONIC, &write_start);
|
|
|
|
// Initialize limits from environment variables (bytes/s)
|
|
char *read_env = getenv("IOTHROTTLE_READ");
|
|
char *write_env = getenv("IOTHROTTLE_WRITE");
|
|
char *debug_env = getenv("IOTHROTTLE_DEBUG");
|
|
|
|
if (debug_env && atoi(debug_env) != 0) {
|
|
debug_mode = 1;
|
|
}
|
|
|
|
if (debug_mode) fprintf(stderr, "init called: default read_limit=%zu, write_limit=%zu\n", read_limit, write_limit);
|
|
|
|
if (read_env) {
|
|
if (debug_mode) fprintf(stderr, "IOTHROTTLE_READ=%s\n", read_env);
|
|
read_limit = strtoull(read_env, NULL, 10);
|
|
} else {
|
|
if (debug_mode) fprintf(stderr, "IOTHROTTLE_READ not set\n");
|
|
}
|
|
|
|
if (write_env) {
|
|
if (debug_mode) fprintf(stderr, "IOTHROTTLE_WRITE=%s\n", write_env);
|
|
write_limit = strtoull(write_env, NULL, 10);
|
|
} else {
|
|
if (debug_mode) fprintf(stderr, "IOTHROTTLE_WRITE not set\n");
|
|
}
|
|
|
|
if (debug_mode) fprintf(stderr, "Final read_limit=%zu, write_limit=%zu\n", read_limit, write_limit);
|
|
}
|