Skip to content

Commit 2aa2900

Browse files
Artur114ProjectsArtur114Projects
committed
Pre release!
1 parent 5619756 commit 2aa2900

18 files changed

Lines changed: 978 additions & 121 deletions

File tree

‎gradle.properties‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
org.gradle.daemon=true
22

33
project_group=com.artur114.bytecodegrab.main
4-
project_archive_name=ByteCodeGrabber
5-
project_version=1.0-SNAPSHOT
4+
project_version=v1.0.0-release
5+
project_archive_name=BCG
66
project_author=Artur114

‎src/main/java/com/artur114/bytecodegrab/agent/Grabber.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import java.lang.instrument.UnmodifiableClassException;
77
import java.lang.reflect.Proxy;
88
import java.util.*;
9+
import java.util.concurrent.atomic.AtomicBoolean;
910

1011
public class Grabber extends Thread {
1112
private final IAgentLogger LOGGER = IAgentLogger.createBestLogger();
@@ -85,13 +86,21 @@ private void sendClassesByteCode(SocketA socket) throws IOException {
8586
try {
8687
socket.send("PROCESS:[" + i + "]");
8788
this.inst.retransformClasses(clazz);
89+
90+
if (!this.validateNext(socket)) {
91+
LOGGER.info("Class retransform aborted");
92+
return;
93+
}
8894
} catch (UnmodifiableClassException e) {
8995
LOGGER.error("Attempt to retransform unmodifiable class, class: {}", clazz);
9096
} catch (Throwable t) {
9197
LOGGER.error("An error occurs while retransform class: {}, skipping...", clazz);
9298
}
9399
}
94100
socket.send("PROCESS:[" + classList.size() + "]");
101+
if (!this.validateNext(socket)) {
102+
return;
103+
}
95104

96105
this.sendByteCode(socket, grabbedClasses);
97106
}
@@ -135,14 +144,24 @@ private int intPropFromMessage(String message, String property) {
135144

136145
private void sendByteCode(SocketA out, Map<String, byte[]> classes) throws IOException {
137146
out.send("CLASS_BYTECODE_LIST_START SIZE:[" + classes.size() + "]");
147+
AtomicBoolean ret = new AtomicBoolean(false);
138148
classes.forEach((key, value) -> {
149+
if (ret.get()) {
150+
return;
151+
}
139152
try {
140153
out.send("CLASS:[" + key + "]");
141154
out.send(value);
155+
156+
ret.set(!this.validateNext(out));
142157
} catch (IOException e) {
143158
e.printStackTrace(System.err);
144159
}
145160
});
161+
if (ret.get()) {
162+
LOGGER.info("Byte code send aborted");
163+
return;
164+
}
146165
out.send("CLASS_BYTECODE_LIST_END");
147166
}
148167

@@ -188,6 +207,20 @@ private boolean isUsefulClassSafe(Class<?> clazz) {
188207
}
189208
}
190209

210+
private boolean validateNext(SocketA connection) throws IOException {
211+
String ret = connection.waitForMessage();
212+
213+
switch (ret) {
214+
case "NEXT":
215+
return true;
216+
case "ABORT":
217+
return false;
218+
default:
219+
LOGGER.warn("Wrong reply {}", ret);
220+
return false;
221+
}
222+
}
223+
191224
private boolean isUsefulClass(Class<?> clazz) {
192225
if (clazz.isPrimitive() || clazz.isArray()) {
193226
return false;

‎src/main/java/com/artur114/bytecodegrab/jcomp/JClassTree.java‎

Lines changed: 53 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
import com.artur114.bytecodegrab.util.AsyncClassTreeBuilder;
44
import com.artur114.bytecodegrab.util.Icons;
5+
import org.slf4j.Logger;
6+
import org.slf4j.LoggerFactory;
57

68
import javax.swing.*;
79

@@ -17,6 +19,7 @@
1719

1820
public class JClassTree extends JPanel {
1921
private static final Color COLOR = new Color(239, 239, 239);
22+
private static final Logger log = LoggerFactory.getLogger(JClassTree.class);
2023
private final Set<String> loadedClasses = new HashSet<>();
2124
private final DefaultTreeModel treeModel;
2225
private AsyncClassTreeBuilder builder;
@@ -38,38 +41,35 @@ public JClassTree() {
3841
@Override
3942
public void mouseReleased(MouseEvent e) {
4043
if (popup != null && e.isPopupTrigger() && SwingUtilities.isRightMouseButton(e)) {
41-
Rectangle r = allRowsRectangle(tree);
42-
int row = tree.getRowForLocation(e.getX(), e.getY());
43-
boolean outOfBounds = !r.contains(e.getPoint()) || row == -1;
44+
int row = rowForY(tree, e.getY());
45+
boolean outOfBounds = row == -1;
4446

4547
if (!outOfBounds) {
4648
tree.requestFocusInWindow();
47-
popup.show(tree, e.getX(), e.getY());
4849
tree.setSelectionRow(row);
50+
popup.show(tree, e.getX(), e.getY());
4951
}
5052
}
5153
}
5254

5355
@Override
5456
public void mouseClicked(MouseEvent e) {
5557
if (popup != null && e.isPopupTrigger() && SwingUtilities.isRightMouseButton(e)) {
56-
Rectangle r = allRowsRectangle(tree);
57-
int row = tree.getRowForLocation(e.getX(), e.getY());
58-
boolean outOfBounds = !r.contains(e.getPoint()) || row == -1;
58+
int row = rowForY(tree, e.getY());
59+
boolean outOfBounds = row == -1;
5960

6061
if (!outOfBounds) {
6162
tree.requestFocusInWindow();
62-
popup.show(tree, e.getX(), e.getY());
6363
tree.setSelectionRow(row);
64+
popup.show(tree, e.getX(), e.getY());
6465
}
6566
}
6667
}
6768

6869
@Override
6970
public void mousePressed(MouseEvent e) {
70-
Rectangle r = allRowsRectangle(tree);
71-
int row = tree.getRowForLocation(e.getX(), e.getY());
72-
boolean outOfBounds = !r.contains(e.getPoint()) || row == -1;
71+
int row = rowForY(tree, e.getY());
72+
boolean outOfBounds = row == -1;
7373
if (SwingUtilities.isLeftMouseButton(e)) {
7474
if (outOfBounds) {
7575
tree.clearSelection();
@@ -80,8 +80,8 @@ public void mousePressed(MouseEvent e) {
8080
}
8181
if (!outOfBounds && popup != null && e.isPopupTrigger() && SwingUtilities.isRightMouseButton(e)) {
8282
tree.requestFocusInWindow();
83-
popup.show(tree, e.getX(), e.getY());
8483
tree.setSelectionRow(row);
84+
popup.show(tree, e.getX(), e.getY());
8585
}
8686
}
8787

@@ -97,6 +97,21 @@ private Rectangle allRowsRectangle(JTree tree) {
9797
}
9898
return ret;
9999
}
100+
101+
private int rowForY(JTree tree, int y) {
102+
if (tree.getRowCount() <= 0) {
103+
return -1;
104+
}
105+
Rectangle ret = tree.getRowBounds(0);
106+
107+
int row = y / ret.height;
108+
109+
if (row > tree.getRowCount()) {
110+
return -1;
111+
}
112+
113+
return row;
114+
}
100115
});
101116

102117
JScrollPane pane = new JScrollPane(this.tree);
@@ -168,7 +183,7 @@ public AsyncClassTreeBuilder updateTree(Set<String> classNames) {
168183
}
169184

170185
public void clear() {
171-
((DefaultMutableTreeNode) this.treeModel.getRoot()).removeAllChildren();
186+
this.treeModel.setRoot(null);
172187
this.loadedClasses.clear();
173188
this.treeModel.reload();
174189
}
@@ -225,6 +240,16 @@ public void expandSelected() {
225240
}
226241
}
227242

243+
public void collapseSelected() {
244+
TreePath[] paths = this.tree.getSelectionPaths();
245+
if (paths == null) {
246+
return;
247+
}
248+
for (TreePath path : paths) {
249+
this.collapseChild(path, (DefaultMutableTreeNode) path.getLastPathComponent());
250+
}
251+
}
252+
228253
public void expandAll() {
229254
for (int i = 0; i != this.tree.getRowCount(); i++) {
230255
this.tree.expandRow(i);
@@ -240,7 +265,12 @@ public void collapseAll() {
240265
}
241266
}
242267

243-
public void expandChild(TreePath path, DefaultMutableTreeNode parent) {
268+
public boolean isSelectedClass() {
269+
TreePath[] paths = this.tree.getSelectionPaths();
270+
return paths != null && paths.length >= 1 && ((DefaultMutableTreeNode) paths[0].getLastPathComponent()).getUserObject() instanceof ClassInfo;
271+
}
272+
273+
private void expandChild(TreePath path, DefaultMutableTreeNode parent) {
244274
this.tree.expandPath(path);
245275

246276
for (DefaultMutableTreeNode node : this.child(parent)) {
@@ -249,6 +279,15 @@ public void expandChild(TreePath path, DefaultMutableTreeNode parent) {
249279
}
250280
}
251281

282+
private void collapseChild(TreePath path, DefaultMutableTreeNode parent) {
283+
for (DefaultMutableTreeNode node : this.child(parent)) {
284+
TreePath newPath = path.pathByAddingChild(node);
285+
this.collapseChild(newPath, node);
286+
}
287+
288+
this.tree.collapsePath(path);
289+
}
290+
252291
@Override
253292
public void setComponentPopupMenu(JPopupMenu popup) {
254293
this.popup = popup;

0 commit comments

Comments
 (0)