feat: Write PoC of sout(access user ip)

easypeel-security · Jan 24, 2024 · 079da0c · 079da0c
1 parent 4c8404b
commit 079da0c
Showing 1 changed file with 21 additions and 5 deletions.
diff --git a/easypeel-method-security/src/main/java/org/epsec/core/MethodBanProcessor.java b/easypeel-method-security/src/main/java/org/epsec/core/MethodBanProcessor.java
@@ -96,24 +96,40 @@ private void generateMethodBanAspect(Element element) {
         .addMember("value", "$S", "@annotation(org.epsec.MethodBan)")
         .build();
 
-    final CodeBlock getMyIpAddressCode = CodeBlock.builder()
-        .addStatement("final String myIpAddress = $T.getHostAddress()",
-            ClassName.get("java.net", "InetAddress"), ClassName.get("java.net", "InetAddress"))
-        .addStatement("System.out.println(myIpAddress)")
+    final MethodSpec getUserIpMethodSpec = MethodSpec.methodBuilder("getUserIp" + System.nanoTime())
+        .returns(String.class)
+        .addModifiers(Modifier.PRIVATE)
+        .addCode(CodeBlock.builder()
+            .addStatement("$T request = (($T) $T.getRequestAttributes()).getRequest()",
+                ClassName.bestGuess("jakarta.servlet.http.HttpServletRequest"),
+                ClassName.bestGuess("org.springframework.web.context.request.ServletRequestAttributes"),
+                ClassName.bestGuess("org.springframework.web.context.request.RequestContextHolder"))
+            .addStatement("$T xForwardedForHeader = $L.getHeader($S)", String.class, "request",
+                "X-Forwarded-For")
+            .beginControlFlow("if ($L == null)", "xForwardedForHeader")
+            .addStatement("return $L.getRemoteAddr()", "request")
+            .endControlFlow()
+            .addStatement("return $L.split($S)[0]", "xForwardedForHeader", ",")
+            .build())
+        .build();
+
+    final CodeBlock soutIp = CodeBlock.builder()
+        .addStatement("$T.out.println($S + $L())", System.class, "IP: ", getUserIpMethodSpec.name)
         .build();
 
     final MethodSpec methodSpec = MethodSpec.methodBuilder("beforeMethodBan" + System.nanoTime())
         .addModifiers(Modifier.PUBLIC)
         .addAnnotation(annotationSpec)
         .addParameter(ClassName.bestGuess(JOIN_POINT.getName()), "joinPoint")
-        .addStatement(getMyIpAddressCode)
+        .addCode(soutIp)
         .build();
 
     final TypeSpec classSpec = TypeSpec.classBuilder("MethodBanAspect" + System.nanoTime())
         .addModifiers(Modifier.PUBLIC)
         .addAnnotation(ClassName.bestGuess(ASPECT.getName()))
         .addAnnotation(ClassName.bestGuess(COMPONENT.getName()))
         .addMethod(methodSpec)
+        .addMethod(getUserIpMethodSpec)
         .build();
 
     try {