srvshim/
srvhook.rs

1extern crate libc;
2extern crate plumber;
3use self::libc::types::os::common::bsd44::{addrinfo, socklen_t, sockaddr};
4use self::libc::{c_char, c_int, size_t, ssize_t};
5use std::collections::{BTreeMap};
6use std::ffi::{CStr};
7use std::mem;
8use std::str::from_utf8;
9use std::sync::{RwLock};
10use std::sync::atomic::{AtomicUsize, Ordering};
11
12use plumber::dynamic::dlsym_next;
13use plumber::util::{sockaddr_to_port_ip,port_ip_to_sa_data};
14use plumber::util::{ip_to_usize,usize_to_ip};
15use plumber::dns::srv_mapper;
16use plumber::hooks::Hook;
17
18pub struct SRVHook {
19    max_ip: AtomicUsize,
20    magic_ip_to_host: RwLock<BTreeMap<[u8;4], String>>,
21    host_to_magic_ip: RwLock<BTreeMap<String, [u8;4]>>,
22    real_connect:
23        unsafe extern "C" fn(c_int, *const sockaddr,
24                             socklen_t) -> c_int,
25    real_getaddrinfo:
26        unsafe extern "C" fn(node: *const c_char,
27                              service: *const c_char,
28                              hints: *const addrinfo,
29                              res: *const *const addrinfo) -> c_int,
30    real_sendto:
31        unsafe extern "C" fn(socket: c_int, msg: *const c_char,
32                             msglen: size_t, flags: c_int,
33                             dest_addr: *mut sockaddr) -> ssize_t,
34}
35
36impl SRVHook {
37    pub unsafe fn new() -> SRVHook {
38        SRVHook{
39            // 250/8 is currently reserved for future use by IANA, so
40            // it should be a safe choice.
41            max_ip: AtomicUsize::new(ip_to_usize([250,0,0,0])),
42            magic_ip_to_host: RwLock::new(BTreeMap::new()),
43            host_to_magic_ip: RwLock::new(BTreeMap::new()),
44            real_getaddrinfo:
45                mem::transmute(dlsym_next("getaddrinfo\0").unwrap()),
46            real_connect:
47                mem::transmute(dlsym_next("connect\0").unwrap()),
48            real_sendto:
49                mem::transmute(dlsym_next("sendto\0").unwrap()),
50        }
51    }
52
53    pub fn set_sockaddr(&self, address: *mut sockaddr) {
54        let (_, ip) = sockaddr_to_port_ip(address);
55        let iph = self.magic_ip_to_host.read().unwrap();
56        iph.get(&ip).map(|h| {
57            srv_mapper(h).map( |(new_port, new_ip)| {
58                unsafe {
59                    // only override host and IP for AF_INET
60                    if (*address).sa_family == 2 {
61                        (*address).sa_data =
62                            port_ip_to_sa_data(new_port, new_ip);
63                    }
64                }
65            }).or_else(|e| {
66                println!("srv-shim: failed to look up SRV record for {}: {}", h, e);
67                Err(e)
68            });
69        });
70    }
71
72    fn next_ip(&self) -> [u8;4] {
73        let id = self.max_ip.fetch_add(1, Ordering::Relaxed);
74        usize_to_ip(id)
75    }
76}
77
78impl Hook for SRVHook {
79    fn connect(&self, socket: c_int, address: *mut sockaddr,
80                   len: socklen_t) -> c_int {
81        self.set_sockaddr(address);
82        unsafe {
83            (self.real_connect)(socket, address, len)
84        }
85    }
86
87    fn sendto(&self, socket: c_int, msg: *const c_char, msglen: size_t,
88                         flags: c_int, dest_addr: *mut sockaddr) -> ssize_t {
89        self.set_sockaddr(dest_addr);
90        unsafe {
91            (self.real_sendto)(socket, msg, msglen, flags, dest_addr)
92        }
93    }
94
95    fn getaddrinfo(&self, node: *const c_char, service: *const c_char,
96                   hints: *const addrinfo, res: *mut *const addrinfo) -> c_int {
97        let c_str = unsafe { CStr::from_ptr(node) };
98        let node_str: String = from_utf8(c_str.to_bytes()).unwrap().to_owned();
99        // Trigger on possible SRV records.
100        if node_str.starts_with("_") {
101            let mut hip = self.host_to_magic_ip.write().unwrap();
102            let ip = match hip.get(&node_str) {
103                Some(ip) => *ip,
104                None => {
105                    let ip = self.next_ip();
106                    let mut iph = self.magic_ip_to_host.write().unwrap();
107                    iph.insert(ip, node_str.clone());
108                    ip
109                },
110            };
111            hip.insert(node_str, ip);
112            unsafe {
113                let sa_buf: *mut sockaddr =
114                    mem::transmute(
115                        libc::malloc(mem::size_of::<sockaddr>() as size_t)
116                    );
117                *sa_buf = sockaddr{
118                    sa_family: 2,
119                    sa_data: port_ip_to_sa_data(8080, ip),
120                };
121
122                let ai_buf: *mut addrinfo =
123                    mem::transmute(
124                        libc::malloc(mem::size_of::<addrinfo>() as size_t)
125                    );
126                *ai_buf = addrinfo{
127                    ai_flags: 0,
128                    ai_family: 2,
129                    ai_socktype: 1,
130                    ai_protocol: 6,
131                    ai_addrlen: 16,
132                    ai_addr: sa_buf,
133                    ai_canonname: 0 as *mut i8,
134                    ai_next: 0 as *mut addrinfo,
135                };
136                *res = ai_buf;
137            }
138            0
139        } else {
140            unsafe {
141                (self.real_getaddrinfo)(node, service, hints, res)
142            }
143        }
144    }
145}