rewrited shred

coreutils/shred.c

@@ -4,36 +4,41 @@
 #include <unistd.h>
 #include <stdlib.h>
 #include <string.h>
+#include <sys/mman.h>
 #include <sys/stat.h>
 
 #define RAND_SOURCE "/dev/urandom"
-#define ZERO_SOURCE "/dev/zero"
 
 unsigned int f_flag;
 unsigned int u_flag;
 unsigned int z_flag;
 unsigned int n_loops;
 
-int shred(int rand_fd, int zero_fd, int fd, unsigned int iter) {
+int shred(int rand_fd, int fd) {
+	/* Get size */
 	off_t size = lseek(fd, 0, SEEK_END);
-	if (size < 0)
+	if (size <= 0)
 		return 1;
 
-	lseek(fd, 0, SEEK_SET);
-
-	char buf[1];
-	while (size) {
-		if (iter == n_loops && z_flag)
-			read(zero_fd, buf, sizeof(buf));
-
-		else
-			read(rand_fd, buf, sizeof(buf));
-
-		write(fd, buf, sizeof(buf));
-		size--;
+	void *buf = mmap(NULL, size, PROT_WRITE, MAP_SHARED, fd, 0);
+	if (buf == MAP_FAILED) {
+		fprintf(stderr, "shred: mmap: %s\n", strerror(errno));
+		return 1;
+	}
+
+	for (unsigned int n = 0; n < n_loops; n++) {
+		read(rand_fd, buf, size);
+		fsync(fd);
+	}
+
+	if (z_flag)
+		memset(buf, 0, size);
+
+	if (munmap(buf, size)) {
+		fprintf(stderr, "shred: munmap: %s\n", strerror(errno));
+		return 1;
 	}
 
-	fsync(fd);
 	return 0;
 }
 
@@ -76,22 +81,14 @@ int main(const int argc, const char **argv) {
 		return 1;
 	}
 
-	int zero_fd = open(ZERO_SOURCE, O_RDONLY);
-	if (zero_fd < 0) {
-		fprintf(stderr, "shred: %s is %s\n", ZERO_SOURCE, strerror(errno));
-		return 1;
-	}
-
 	for (; i < argc; i++) {
-		int fd = open(argv[i], O_WRONLY);
+		int fd = open(argv[i], O_RDWR);
 		if (fd < 0) {
 			fprintf(stderr, "shred: %s is %s\n", argv[i], strerror(errno));
-			continue;
+			return 1;
 		}
 
-		for (unsigned int n = 0; n <= n_loops; n++)
-			shred(rand_fd, zero_fd, fd, n);
-
+		shred(rand_fd, fd);
 		close(fd);
 
 		if (u_flag)
@@ -104,6 +101,5 @@ int main(const int argc, const char **argv) {
 	}
 
 	close(rand_fd);
-	close(zero_fd);
 	return 0;
 }