Fix the bug if the file size is less than 1024.

tinystruct · Apr 19, 2024 · 0052a9d · 0052a9d
1 parent ca96b63
commit 0052a9d
Showing 1 changed file with 17 additions and 18 deletions.
diff --git a/src/main/java/custom/application/v1/smalltalk.java b/src/main/java/custom/application/v1/smalltalk.java
@@ -767,25 +767,22 @@ public String upload() throws ApplicationException {
                  final BufferedOutputStream bout = new BufferedOutputStream(out);
                  final BufferedInputStream bs = new BufferedInputStream(new ByteArrayInputStream(file.get()));
             ) {
-                final byte[] bytes = new byte[1024];
-                byte[] keys = meetingCode.toString().getBytes(StandardCharsets.UTF_8);
-                int read;
-                while ((read = bs.read(bytes)) != -1) {
-                    for (int i = 0; i < keys.length; i++) {
-                        bytes[i] = (byte) (bytes[i] ^ keys[i]);
+                byte[] keyBytes = keyBytes(meetingCode.toString()); // Generate key bytes
+                int bytesRead;
+                byte[] buffer = new byte[1024];
+                while ((bytesRead = bs.read(buffer)) != -1) {
+                    for (int i = 0; i < bytesRead; i++) {
+                        buffer[i] = (byte) (buffer[i] ^ keyBytes[i % keyBytes.length]); // XOR operation
                     }
-                    bout.write(bytes, 0, read);
+                    bout.write(buffer, 0, bytesRead);
                 }
 
-                bout.close();
-                bs.close();
-
                 builders.add(builder);
                 System.out.printf("File %s being uploaded to %s%n", file.getFilename(), path);
             } catch (FileNotFoundException e) {
-                throw new ApplicationException(e.getMessage(), e);
+                throw new ApplicationException("File not found: " + e.getMessage(), e);
             } catch (IOException e) {
-                throw new ApplicationException(e.getMessage(), e);
+                throw new ApplicationException("Error uploading file: " + e.getMessage(), e);
             }
         }
 
@@ -818,20 +815,22 @@ public byte[] download(String fileName, boolean encoded) throws ApplicationExcep
 
             arr = Files.readAllBytes(path);
             if (encoded) {
-                byte[] keys = meetingCode.toString().getBytes(StandardCharsets.UTF_8);
-                for (int i = 0; i < arr.length; i = i + 1024) {
-                    for (int j = 0; j < keys.length; j++) {
-                        arr[i + j] = (byte) (arr[i + j] ^ keys[j]);
-                    }
+                byte[] keyBytes = keyBytes(meetingCode.toString()); // Generate key bytes
+                for (int i = 0; i < arr.length; i++) {
+                    arr[i] = (byte) (arr[i] ^ keyBytes[i % keyBytes.length]);
                 }
             }
         } catch (IOException e) {
-            e.printStackTrace();
+            throw new ApplicationException("Error reading the file: " + e.getMessage(), e);
         }
 
         return arr;
     }
 
+    private byte[] keyBytes(String meetingCode) {
+        return meetingCode.getBytes(StandardCharsets.UTF_8);
+    }
+
     @Action("files")
     public byte[] download(String fileName) throws ApplicationException {
         return this.download(fileName, true);