Merge pull request #29 from jpsamaroo/jps/stacks-on-stacks

Add stack helpers and more

Project.toml

@@ -1,7 +1,7 @@
 name = "BPFnative"
 uuid = "b6338580-32ea-11e9-1791-33a79977d8c4"
 authors = ["Julian P Samaroo <[email protected]>"]
-version = "0.1.4"
+version = "0.1.5"
 
 [deps]
 CBinding = "d43a6710-96b8-4a2d-833c-c424785e5374"

src/BPFnative.jl

@@ -66,6 +66,10 @@ end
 # Runtime API
 module RT
 import ..API
+using ..LLVM
+using ..LLVM.Interop
+import Core: LLVMPtr
+include("runtime/maps_core.jl")
 include("runtime/bpfcall.jl")
 include("runtime/maps.jl")
 include("runtime/buffers.jl")
@@ -79,6 +83,7 @@ import ..API
 include("host/syscall.jl")
 include("host/maps.jl")
 include("host/socket.jl")
+include("host/kallsyms.jl")
 end
 
 # Compiler

src/common.jl

@@ -255,6 +255,8 @@ end
     sk_assign
 end
 
+const PERF_MAX_STACK_DEPTH = 127
+
 # Kernel structures
 
 if has_vmlinux

src/host/kallsyms.jl

@@ -0,0 +1,32 @@
+function find_ksym(addr::UInt64)
+    start, stop = UInt64(0), addr
+    rgx = r"([0-9abcdef]*) ([a-zA-Z]) ([0-9a-zA-Z_\-]*)"
+    last_addr = start
+    last_kind = "?"
+    last_name = ""
+    for line in readlines(open("/proc/kallsyms", "r"))
+        m = match(rgx, line)
+        @assert m !== nothing
+        start_addr, kind, name = m.captures
+        start_addr = parse(UInt64, "0x"*start_addr)
+        if start_addr > stop
+            return last_addr, last_kind, last_name
+        elseif start_addr == stop
+            return addr, kind, name
+        end
+        last_addr = addr
+        last_kind = kind
+        last_name = name
+    end
+end
+function stack_to_string(nt::NTuple{N,UInt64}) where N
+    iob = IOBuffer()
+    for i in 1:N
+        addr = nt[i]
+        if addr == UInt64(0)
+            break
+        end
+        println(iob, "$(find_ksym(addr)[3])")
+    end
+    String(take!(iob))
+end

src/host/maps.jl

@@ -15,6 +15,7 @@ struct HashMap{K,V} <: AbstractHashMap{K,V}
     fd::Cint
 end
 maptype_to_jltype(::Val{API.BPF_MAP_TYPE_HASH}) = HashMap
+maptype_to_jltype(::Val{API.BPF_MAP_TYPE_STACK_TRACE}) = HashMap
 struct ArrayMap{K,V} <: AbstractArrayMap{K,V}
     fd::Cint
 end
@@ -45,8 +46,25 @@ struct map_access_elem_attr
     flags::UInt64
 end
 
+memset!(ptr::Ptr{T}) where T =
+    ccall(:memset, Cvoid,
+          (Ptr{T}, UInt8, UInt64),
+          ptr, UInt8(0), sizeof(T))
+"Creates a Ref{T} that's been zero-initialized before being stored. Necessary
+to ensure that struct padding bytes are zeroed."
+function ZeroInitRef(T, val; set=true)
+    ref = Ref{T}()
+    memset!(Base.unsafe_convert(Ptr{T}, ref))
+    if set
+        ref[] = val
+    end
+    ref
+end
+ZeroInitRef(val::T) where T = ZeroInitRef(T, val)
+ZeroInitRef(::Type{T}) where T = ZeroInitRef(T, nothing; set=false)
+
 function Base.getindex(map::AbstractHashMap{K,V}, idx) where {K,V}
-    key = Ref{K}(idx)
+    key = ZeroInitRef(K, idx)
     value = Ref{V}()
     key_ptr = Base.unsafe_convert(Ptr{K}, key)
     value_ptr = Base.unsafe_convert(Ptr{V}, value)
@@ -61,7 +79,7 @@ function Base.getindex(map::AbstractHashMap{K,V}, idx) where {K,V}
     value[]
 end
 function Base.getindex(map::AbstractArrayMap{K,V}, idx) where {K,V}
-    key = Ref{K}(idx-1)
+    key = ZeroInitRef(K, idx-1)
     value = Ref{V}()
     key_ptr = Base.unsafe_convert(Ptr{K}, key)
     value_ptr = Base.unsafe_convert(Ptr{V}, value)
@@ -77,8 +95,8 @@ function Base.getindex(map::AbstractArrayMap{K,V}, idx) where {K,V}
 end
 
 function Base.setindex!(map::AbstractHashMap{K,V}, value::U, idx) where {K,V,U}
-    key_ref = Ref{K}(convert(K,idx))
-    value_ref = Ref{V}(convert(V,value))
+    key_ref = ZeroInitRef(convert(K,idx))
+    value_ref = ZeroInitRef(convert(V,value))
     key_ptr = Base.unsafe_convert(Ptr{K}, key_ref)
     value_ptr = Base.unsafe_convert(Ptr{V}, value_ref)
     attr = Ref(map_access_elem_attr(map.fd,
@@ -92,8 +110,8 @@ function Base.setindex!(map::AbstractHashMap{K,V}, value::U, idx) where {K,V,U}
     value
 end
 function Base.setindex!(map::AbstractArrayMap{K,V}, value::U, idx) where {K,V,U}
-    key_ref = Ref{K}(convert(K,idx-1))
-    value_ref = Ref{V}(convert(V,value))
+    key_ref = ZeroInitRef(convert(K,idx-1))
+    value_ref = ZeroInitRef(convert(V,value))
     key_ptr = Base.unsafe_convert(Ptr{K}, key_ref)
     value_ptr = Base.unsafe_convert(Ptr{V}, value_ref)
     attr = Ref(map_access_elem_attr(map.fd,
@@ -108,7 +126,7 @@ function Base.setindex!(map::AbstractArrayMap{K,V}, value::U, idx) where {K,V,U}
 end
 
 function Base.delete!(map::AbstractHashMap{K,V}, idx) where {K,V}
-    key = Ref{K}(idx)
+    key = ZeroInitRef(K, idx)
     key_ptr = Base.unsafe_convert(Ptr{K}, key)
     attr = Ref(map_access_elem_attr(map.fd,
                                     key_ptr,
@@ -121,7 +139,7 @@ function Base.delete!(map::AbstractHashMap{K,V}, idx) where {K,V}
 end
 
 function Base.haskey(map::HostMap{K,V}, idx) where {K,V}
-    key = Ref{K}(idx)
+    key = ZeroInitRef(K, idx)
     value = Ref{V}()
     key_ptr = Base.unsafe_convert(Ptr{K}, key)
     value_ptr = Base.unsafe_convert(Ptr{V}, value)
@@ -134,4 +152,47 @@ function Base.haskey(map::HostMap{K,V}, idx) where {K,V}
     end
 end
 
-# TODO: keys(map) using BPF_MAP_GET_NEXT_KEY
+function nextkey(map::HostMap{K,V}, idx) where {K,V}
+    key = ZeroInitRef(K, idx)
+    nkey = Ref{K}()
+    key_ptr = Base.unsafe_convert(Ptr{K}, key)
+    nkey_ptr = Base.unsafe_convert(Ptr{K}, nkey)
+    attr = Ref(map_access_elem_attr(map.fd,
+                                    key_ptr,
+                                    nkey_ptr,
+                                    0))
+    ret = GC.@preserve key nkey begin
+        bpf(API.BPF_MAP_GET_NEXT_KEY, attr)
+    end
+    if (ret == -1) && (Libc.errno() == Libc.ENOENT)
+        return nothing
+    end
+    ret >= 0 || Base.systemerror(ret)
+    nkey[]
+end
+
+struct HostMapKeySet{K,V,H<:HostMap{K,V}}
+    map::H
+end
+Base.keys(map::H) where {K,V,H<:HostMap{K,V}} = HostMapKeySet{K,V,H}(map)
+Base.IteratorSize(::Type{<:HostMapKeySet}) = Base.SizeUnknown()
+Base.eltype(::Type{HostMapKeySet{K,V,H}}) where {K,V,H} = K
+function Base.iterate(hmks::HostMapKeySet{K,V,H}) where {K,V,H}
+    fakekey_ref = ZeroInitRef(K)
+    realkey = if haskey(hmks.map, fakekey_ref[])
+        fakekey_ref[]
+    else
+        nextkey(hmks.map, fakekey_ref[])
+    end
+    realkey === nothing && return nothing # empty
+    return realkey, (realkey, realkey)
+end
+function Base.iterate(hmks::HostMapKeySet{K,V,H}, (lastkey, initial)) where {K,V,H}
+    lastkey === nothing && return nothing
+    nkey = nextkey(hmks.map, lastkey)
+    if nkey === nothing || nkey == lastkey
+        return nothing
+    end
+    return nkey, (nkey == initial ? nothing : nkey, initial)
+end
+Base.length(map::HostMap) = length(keys(map))

src/runtime/bpfcall.jl

@@ -3,21 +3,52 @@ export bpfcall
 @generated function _bpfcall(::Val{cmd}, ::Type{rettype}, ::Type{argtypes}, args::Vararg{Any}) where {cmd,rettype,argtypes}
     JuliaContext() do ctx
         T_ret = convert(LLVMType, rettype, ctx)
-        T_args = map(x->convert(LLVMType, x, ctx), argtypes.parameters)
+        T_jlptr = convert(LLVMType, Ptr{Cvoid}, ctx)
+        T_ptr_i8 = LLVM.PointerType(LLVM.Int8Type(ctx))
 
-        llvm_f, _ = create_function(T_ret, LLVMType[T_args...])
+        outer_args = filter(arg->!(arg <: RTMap), args)
+        T_outer_args = LLVMType[convert(LLVMType, arg, ctx) for arg in outer_args]
+
+        llvm_f, _ = create_function(T_ret, T_outer_args)
         mod = LLVM.parent(llvm_f)
 
         Builder(ctx) do builder
             entry = BasicBlock(llvm_f, "entry", ctx)
             position!(builder, entry)
-            ft = LLVM.FunctionType(T_ret, LLVMType[T_args...])
+
+            inner_args = LLVM.Value[]
+            T_inner_args = LLVMType[]
+            outer_args_ex = Expr[]
+            jlidx = 1
+            pidx = 1
+            for (idx, arg) in enumerate(argtypes.parameters)
+                if arg <: RTMap
+                    map_gv = _genmap!(mod, arg, ctx)
+                    push!(inner_args, map_gv)
+                    push!(T_inner_args, llvmtype(map_gv))
+                else
+                    llvm_arg = parameters(llvm_f)[pidx]
+                    T_arg = T_outer_args[pidx]
+                    if arg <: Ptr
+                        llvm_arg = inttoptr!(builder, llvm_arg, T_ptr_i8)
+                        T_arg = llvmtype(llvm_arg)
+                    end
+                    push!(inner_args, llvm_arg)
+                    push!(T_inner_args, T_arg)
+                    push!(outer_args_ex, Expr(:ref, :args, jlidx))
+                    pidx += 1
+                end
+                jlidx += 1
+            end
+
+            ft = LLVM.FunctionType(T_ret, T_inner_args)
             ftp = LLVM.PointerType(ft)
             f = inttoptr!(builder, ConstantInt(cmd, ctx), ftp)
-            value = call!(builder, f, LLVM.Value[parameters(llvm_f)...])
+            value = call!(builder, f, inner_args)
             ret!(builder, value)
         end
-        call_function(llvm_f, rettype, Base.to_tuple_type(args), :((args...,)))
+        outer_args_ex = Expr(:tuple, outer_args_ex...)
+        call_function(llvm_f, rettype, Base.to_tuple_type(outer_args), outer_args_ex)
     end
 end
 @inline bpfcall(cmd::API.BPFHelper, RT, AT, args...) =

src/runtime/helpers.jl

@@ -1,6 +1,7 @@
 # BPF helpers
 
 const ptr_sk_buff = API.pointertype(API.sk_buff)
+const ptr_task_struct = Ptr{Cvoid} #API.pointertype(API.task_struct)
 
 bpfconvert(x) = x
 bpfconvert(x::AbstractBuffer) = pointer(x)
@@ -54,4 +55,19 @@ end
 @inline get_current_comm(buf::AbstractSizedBuffer) =
     bpfcall(API.get_current_comm, Clong, Tuple{BufPtr, UInt32}, pointer(buf), length(buf))
 
+@inline get_stackid(ctx::T, map::M, flags::Integer) where {T,M<:RTMap} =
+    bpfcall(API.get_stackid, Clong, Tuple{T, M, UInt64}, ctx, map, unsafe_trunc(UInt64,flags))
+@inline function get_current_task()
+    res = bpfcall(API.get_current_task, UInt64)
+    if res > 0
+        unsafe_load(reinterpret(ptr_task_struct, res))
+    else
+        nothing
+    end
+end
+@inline get_stack(ctx::T, buf::AbstractSizedBuffer, flags::UInt64) where {T} =
+    bpfcall(API.get_stack, Clong, Tuple{T, BufPtr, UInt32, UInt64}, ctx, pointer(buf), length(buf), flags)
+@inline get_task_stack(ctx::ptr_task_struct, buf::AbstractSizedBuffer, flags::UInt64) where {T} =
+    bpfcall(API.get_task_stack, Clong, Tuple{ptr_task_struct, BufPtr, UInt32, UInt64}, ctx, pointer(buf), length(buf), flags)
+
 # TODO: The rest!