--- linux-2.6.35.1/net/ipv4/tcp_probe.c	2010-08-15 22:40:59.306549344 -0400
+++ linux-2.6.35-soheil/net/ipv4/tcp_probe.c	2010-08-15 22:41:34.506558165 -0400
@@ -3,6 +3,8 @@
  *
  * The idea for this came from Werner Almesberger's umlsim
  * Copyright (C) 2004, Stephen Hemminger <shemminger@osdl.org>
+ * Copyright (C) 2010, Soheil Hassas Yeganeh <soheil@cs.toronto.edu> 
+ *      Removing kprobes dependency
  *
  * This program is free software; you can redistribute it and/or modify
  * it under the terms of the GNU General Public License as published by
@@ -124,18 +126,12 @@
 
 		wake_up(&tcp_probe.wait);
 	}
-
+#ifdef CONFIG_KPROBES
 	jprobe_return();
+#endif
 	return 0;
 }
 
-static struct jprobe tcp_jprobe = {
-	.kp = {
-		.symbol_name	= "tcp_rcv_established",
-	},
-	.entry	= jtcp_rcv_established,
-};
-
 static int tcpprobe_open(struct inode * inode, struct file * file)
 {
 	/* Reset (empty) log */
@@ -170,6 +166,7 @@
 	int error = 0;
 	size_t cnt = 0;
 
+
 	if (!buf)
 		return -EINVAL;
 
@@ -177,6 +174,7 @@
 		char tbuf[128];
 		int width;
 
+    
 		/* Wait for data in buffer */
 		error = wait_event_interruptible(tcp_probe.wait,
 						 tcp_probe_used() > 0);
@@ -216,6 +214,130 @@
 	.read    = tcpprobe_read,
 };
 
+#ifdef CONFIG_KPROBES
+
+static struct jprobe tcp_jprobe = {
+    .kp = {
+        .symbol_name	= "tcp_rcv_established",
+    },
+    .entry	= jtcp_rcv_established,
+};
+
+#define register_tcpprobe_hook(ret) ret = register_jprobe(&tcp_jprobe) 
+#define unregister_tcpprobe_hook() unregister_jprobe(&tcp_jprobe)
+#else
+static int (*orig_tcp_v4_do_rcv) (struct sock* sk, struct sk_buff* skb);
+
+int tcpprobe_v4_do_rcv(struct sock *sk, struct sk_buff *skb)
+{
+    int ret;
+    try_module_get(THIS_MODULE);
+    ret = (*orig_tcp_v4_do_rcv)(sk, skb); 
+    if(sk->sk_state == TCP_ESTABLISHED) {
+        jtcp_rcv_established(sk, skb, tcp_hdr(skb), skb->len);
+    }
+    module_put(THIS_MODULE);
+    return ret;
+}
+
+static inline int is_empty_ebucket(int bucket)
+{
+    return hlist_nulls_empty(&tcp_hashinfo.ehash[bucket].chain) &&
+        hlist_nulls_empty(&tcp_hashinfo.ehash[bucket].twchain);
+}
+
+
+static inline void register_in_tcp_prot(void) 
+{
+    int bucket = 0;
+    orig_tcp_v4_do_rcv = tcp_prot.backlog_rcv;
+    tcp_prot.backlog_rcv = tcpprobe_v4_do_rcv;
+    while (bucket <= tcp_hashinfo.ehash_mask) {
+        if (!is_empty_ebucket(bucket)) {    
+            struct inet_ehash_bucket* ehash_bucket = inet_ehash_bucket(&tcp_hashinfo, bucket);
+            spinlock_t *lock = inet_ehash_lockp(&tcp_hashinfo, bucket);
+            struct sock *sk = sk_nulls_head(&ehash_bucket->chain);
+            struct hlist_nulls_node* node; 
+            spin_lock(lock);
+            sk_nulls_for_each(sk, node, &ehash_bucket->chain) {
+                if(sk->sk_prot == &tcp_prot) {
+                    sk->sk_backlog_rcv = tcpprobe_v4_do_rcv;
+                }
+            }
+            sk = sk_nulls_head(&ehash_bucket->twchain);
+            sk_nulls_for_each(sk, node, &ehash_bucket->twchain) {
+                if(sk->sk_prot == &tcp_prot) {
+                    sk->sk_backlog_rcv = tcpprobe_v4_do_rcv;
+                } 
+            }
+            spin_unlock(lock);
+        }
+        bucket++;
+    }
+
+    for(bucket = 0; bucket < INET_LHTABLE_SIZE; bucket++) {
+        struct inet_listen_hashbucket* lhash_bucket = &tcp_hashinfo.listening_hash[bucket];
+        spinlock_t *lock = &lhash_bucket->lock;
+        struct sock *sk = sk_nulls_head(&lhash_bucket->head);
+        struct hlist_nulls_node* node; 
+        spin_lock(lock);
+        sk_nulls_for_each(sk, node, &lhash_bucket->head) {
+            if(sk->sk_prot == &tcp_prot) {
+                sk->sk_backlog_rcv = tcpprobe_v4_do_rcv;
+            } 
+        }
+        spin_unlock(lock);
+    }
+
+}
+
+static inline void unregister_in_tcp_prot(void)
+{
+    int bucket = 0;
+    tcp_prot.backlog_rcv = orig_tcp_v4_do_rcv;
+    while(bucket <= tcp_hashinfo.ehash_mask) {
+        if(!is_empty_ebucket(bucket)) {    
+            struct inet_ehash_bucket* ehash_bucket = inet_ehash_bucket(&tcp_hashinfo, bucket);
+            spinlock_t *lock = inet_ehash_lockp(&tcp_hashinfo, bucket);
+            struct sock *sk = sk_nulls_head(&ehash_bucket->chain);
+            struct hlist_nulls_node* node; 
+            spin_lock(lock);
+            sk_nulls_for_each(sk, node, &ehash_bucket->chain) {
+                if(sk->sk_prot == &tcp_prot) {
+                    sk->sk_backlog_rcv = orig_tcp_v4_do_rcv;
+                } 
+            }
+            sk = sk_nulls_head(&ehash_bucket->twchain);
+            sk_nulls_for_each(sk, node, &ehash_bucket->twchain) {
+                if(sk->sk_prot == &tcp_prot) {
+                    sk->sk_backlog_rcv = orig_tcp_v4_do_rcv;
+                }
+            }
+            spin_unlock(lock);
+        }
+        bucket++;
+    }
+
+    for(bucket = 0; bucket < INET_LHTABLE_SIZE; bucket++) {
+        struct inet_listen_hashbucket* lhash_bucket = &tcp_hashinfo.listening_hash[bucket];
+        spinlock_t *lock = &lhash_bucket->lock;
+        struct sock *sk = sk_nulls_head(&lhash_bucket->head);
+        struct hlist_nulls_node* node; 
+        spin_lock(lock);
+        sk_nulls_for_each(sk, node, &lhash_bucket->head) {
+            if(sk->sk_prot == &tcp_prot) {
+                sk->sk_backlog_rcv = orig_tcp_v4_do_rcv;
+            }
+        }
+        spin_unlock(lock);
+    }
+
+}
+#define register_tcpprobe_hook(ret)  { ret = 0; register_in_tcp_prot();}
+#define unregister_tcpprobe_hook()  { unregister_in_tcp_prot(); } 
+#endif
+
+
 static __init int tcpprobe_init(void)
 {
 	int ret = -ENOMEM;
@@ -234,7 +356,7 @@
 	if (!proc_net_fops_create(&init_net, procname, S_IRUSR, &tcpprobe_fops))
 		goto err0;
 
-	ret = register_jprobe(&tcp_jprobe);
+    register_tcpprobe_hook(ret);
 	if (ret)
 		goto err1;
 
@@ -250,8 +372,10 @@
 
 static __exit void tcpprobe_exit(void)
 {
+    spin_lock(&tcp_probe.lock);
 	proc_net_remove(&init_net, procname);
-	unregister_jprobe(&tcp_jprobe);
+    unregister_tcpprobe_hook();
 	kfree(tcp_probe.log);
+    spin_unlock(&tcp_probe.lock);
 }
 module_exit(tcpprobe_exit);
