Berkeley Packet Filter (BPF) assembler. https://godoc.org/maze.io/x/bpf
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

assembler.go 11KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443
  1. //go:generate antlr -package parser -Dlanguage=Go -visitor internal/parser/Assembler.g4
  2. // Copyright (c) 2019 Wijnand Modderman-Lenstra. All rights reserved.
  3. // Use of this source code is governed by a BSD-style
  4. // license that can be found in the LICENSE file.
  5. package bpf
  6. import (
  7. "errors"
  8. "fmt"
  9. "github.com/antlr/antlr4/runtime/Go/antlr"
  10. "golang.org/x/net/bpf"
  11. "maze.io/x/bpf/internal/parser"
  12. )
  13. // Assemble BPF instructions from source.
  14. func Assemble(source string) (Program, error) {
  15. var (
  16. input = antlr.NewInputStream(source)
  17. lexer = parser.NewAssemblerLexer(input)
  18. stream = antlr.NewCommonTokenStream(lexer, 0)
  19. parse = parser.NewAssemblerParser(stream)
  20. )
  21. parse.BuildParseTrees = true
  22. asm := new(assembler)
  23. result := asm.Visit(parse.Program()).([]interface{})
  24. if result[1] == nil {
  25. return result[0].([]bpf.Instruction), nil
  26. }
  27. return nil, result[1].(error)
  28. }
  29. type assembler struct {
  30. *antlr.BaseParseTreeVisitor
  31. }
  32. func (asm *assembler) Visit(tree antlr.ParseTree) interface{} {
  33. return tree.Accept(asm)
  34. }
  35. func (asm *assembler) VisitChildren(node antlr.RuleNode) interface{} {
  36. var results []interface{}
  37. for _, child := range node.GetChildren() {
  38. switch child := child.(type) {
  39. case *antlr.ErrorNodeImpl:
  40. results = append(results, fmt.Errorf("invalid token %q", child.GetText()))
  41. default:
  42. instruction := child.(antlr.ParseTree).Accept(asm)
  43. if instruction != nil {
  44. results = append(results, instruction)
  45. }
  46. }
  47. }
  48. return results
  49. }
  50. func (asm *assembler) VisitProgram(ctx *parser.ProgramContext) interface{} {
  51. program := []interface{}{}
  52. labelAddresses := map[string]uint32{}
  53. results := asm.VisitChildren(ctx)
  54. for _, result := range results.([]interface{}) {
  55. switch result := result.(type) {
  56. case string:
  57. labelAddresses[result] = uint32(len(program))
  58. //log.Printf("label %q at %d", result, labelAddresses[result])
  59. case []interface{}:
  60. for _, word := range result {
  61. switch word := word.(type) {
  62. case string:
  63. labelAddresses[word] = uint32(len(program))
  64. default:
  65. program = append(program, word)
  66. }
  67. }
  68. case error:
  69. return []interface{}{nil, result}
  70. default:
  71. program = append(program, result)
  72. }
  73. }
  74. resolveLabelAddress := func(i int, label string) (skip uint8, err error) {
  75. if addr, ok := labelAddresses[label]; ok {
  76. // log.Printf("label %q at %d -> %d", label, addr, addr-uint32(i)-1)
  77. if addr < uint32(i) {
  78. return 0, fmt.Errorf("instruction %d: jump to negative offset label %q", i, label)
  79. }
  80. return uint8(addr - uint32(i) - 1), nil
  81. }
  82. return 0, fmt.Errorf("instruction %d: jump to unresolved label %q", i, label)
  83. }
  84. var instructions []bpf.Instruction
  85. for i, word := range program {
  86. // log.Printf("program: %T: %+v", word, word)
  87. switch word := word.(type) {
  88. case bpf.Instruction:
  89. instructions = append(instructions, word)
  90. case conditionalJump:
  91. switch skip := word.skipTrue.(type) {
  92. case uint8:
  93. word.Instruction.SkipTrue = skip
  94. case string:
  95. addr, err := resolveLabelAddress(i, skip)
  96. if err != nil {
  97. return []interface{}{nil, err}
  98. }
  99. word.Instruction.SkipFalse = addr
  100. default:
  101. return []interface{}{nil, fmt.Errorf("instruction %d: invalid jump true %T", i, skip)}
  102. }
  103. switch skip := word.skipFalse.(type) {
  104. case uint8:
  105. word.Instruction.SkipFalse = skip
  106. case string:
  107. addr, err := resolveLabelAddress(i, skip)
  108. if err != nil {
  109. return []interface{}{nil, err}
  110. }
  111. word.Instruction.SkipFalse = addr
  112. }
  113. instructions = append(instructions, word.Instruction)
  114. case conditionalJumpX:
  115. switch skip := word.skipTrue.(type) {
  116. case uint8:
  117. word.Instruction.SkipTrue = skip
  118. case string:
  119. addr, err := resolveLabelAddress(i, skip)
  120. if err != nil {
  121. return []interface{}{nil, err}
  122. }
  123. word.Instruction.SkipFalse = addr
  124. }
  125. switch skip := word.skipFalse.(type) {
  126. case uint8:
  127. word.Instruction.SkipFalse = skip
  128. case string:
  129. addr, err := resolveLabelAddress(i, skip)
  130. if err != nil {
  131. return []interface{}{nil, err}
  132. }
  133. word.Instruction.SkipFalse = addr
  134. }
  135. instructions = append(instructions, word.Instruction)
  136. default:
  137. return []interface{}{nil, fmt.Errorf("instruction %d: unknown type %T", i, word)}
  138. }
  139. }
  140. return []interface{}{instructions, nil}
  141. }
  142. func (asm *assembler) VisitLabelDefinition(ctx *parser.LabelDefinitionContext) interface{} {
  143. return ctx.IDENTIFIER().GetText()
  144. }
  145. func (asm *assembler) VisitLabel(ctx *parser.LabelContext) interface{} {
  146. return ctx.IDENTIFIER().GetText()
  147. }
  148. func (asm *assembler) VisitComment(ctx *parser.CommentContext) interface{} {
  149. return asm.VisitChildren(ctx)
  150. }
  151. func (asm *assembler) VisitInstruction(ctx *parser.InstructionContext) interface{} {
  152. switch {
  153. case ctx.AluOperation() != nil:
  154. return asm.Visit(ctx.AluOperation())
  155. case ctx.JumpOperation() != nil:
  156. return asm.Visit(ctx.JumpOperation())
  157. case ctx.JumpConditionalOperation() != nil:
  158. return asm.Visit(ctx.JumpConditionalOperation())
  159. case ctx.LoadOperation() != nil:
  160. return asm.Visit(ctx.LoadOperation())
  161. case ctx.StoreOperation() != nil:
  162. return asm.Visit(ctx.StoreOperation())
  163. case ctx.SimpleOperation() != nil:
  164. return asm.Visit(ctx.SimpleOperation())
  165. case ctx.ReturnOperation() != nil:
  166. return asm.Visit(ctx.ReturnOperation())
  167. default:
  168. return fmt.Errorf("unknown instruction: %#+v", ctx)
  169. }
  170. }
  171. func (asm *assembler) VisitAluOperation(ctx *parser.AluOperationContext) interface{} {
  172. var (
  173. opcode = asm.Visit(ctx.AluOperator()).(bpf.ALUOp)
  174. )
  175. switch {
  176. case ctx.LiteralNumber() != nil:
  177. return bpf.ALUOpConstant{Op: opcode, Val: asm.Visit(ctx.LiteralNumber()).(uint32)}
  178. case ctx.RegisterX() != nil:
  179. return bpf.ALUOpX{Op: opcode}
  180. default:
  181. return nil
  182. }
  183. }
  184. func (asm *assembler) VisitAluOperator(ctx *parser.AluOperatorContext) interface{} {
  185. return parseAluOpcode(ctx)
  186. }
  187. func (asm *assembler) VisitJumpOperation(ctx *parser.JumpOperationContext) interface{} {
  188. return asm.VisitChildren(ctx)
  189. }
  190. func (asm *assembler) VisitJumpOperator(ctx *parser.JumpOperatorContext) interface{} {
  191. return asm.VisitChildren(ctx)
  192. }
  193. type conditionalJump struct {
  194. Instruction bpf.JumpIf
  195. skipTrue interface{}
  196. skipFalse interface{}
  197. }
  198. type conditionalJumpX struct {
  199. Instruction bpf.JumpIfX
  200. skipTrue interface{}
  201. skipFalse interface{}
  202. }
  203. func (asm *assembler) VisitJumpConditionalOperation(ctx *parser.JumpConditionalOperationContext) interface{} {
  204. var (
  205. test = asm.Visit(ctx.JumpConditionalOperator()).(bpf.JumpTest)
  206. ifTrue = asm.Visit(ctx.IfTrue())
  207. ifFalse interface{}
  208. )
  209. if ctx.IfFalse() != nil {
  210. ifFalse = asm.Visit(ctx.IfFalse())
  211. }
  212. switch {
  213. case ctx.LiteralNumber() != nil:
  214. return conditionalJump{
  215. Instruction: bpf.JumpIf{
  216. Cond: test,
  217. Val: asm.Visit(ctx.LiteralNumber()).(uint32),
  218. },
  219. skipTrue: ifTrue,
  220. skipFalse: ifFalse,
  221. }
  222. case ctx.RegisterX() != nil:
  223. return conditionalJumpX{
  224. Instruction: bpf.JumpIfX{
  225. Cond: test,
  226. },
  227. skipTrue: ifTrue,
  228. skipFalse: ifFalse,
  229. }
  230. default:
  231. return errors.New("invalid conditional jump")
  232. }
  233. }
  234. func (asm *assembler) VisitJumpConditionalOperator(ctx *parser.JumpConditionalOperatorContext) interface{} {
  235. return parseJumpCondition(ctx)
  236. }
  237. func (asm *assembler) VisitIfTrue(ctx *parser.IfTrueContext) interface{} {
  238. switch {
  239. case ctx.Label() != nil:
  240. return asm.Visit(ctx.Label())
  241. default:
  242. return uint8(asm.Visit(ctx.Number()).(uint32))
  243. }
  244. }
  245. func (asm *assembler) VisitIfFalse(ctx *parser.IfFalseContext) interface{} {
  246. switch {
  247. case ctx.Label() != nil:
  248. return asm.Visit(ctx.Label())
  249. default:
  250. return uint8(asm.Visit(ctx.Number()).(uint32))
  251. }
  252. }
  253. func (asm *assembler) VisitLoadOperation(ctx *parser.LoadOperationContext) interface{} {
  254. target := asm.Visit(ctx.LoadOperator()).(dst)
  255. switch {
  256. case ctx.AbsoluteNumber() != nil:
  257. var (
  258. number = asm.Visit(ctx.AbsoluteNumber()).(uint32)
  259. instruction = bpf.LoadConstant{Val: number}
  260. )
  261. switch target {
  262. case regA:
  263. instruction.Dst = bpf.RegA
  264. case regX:
  265. instruction.Dst = bpf.RegX
  266. }
  267. return instruction
  268. case ctx.RegisterR() != nil:
  269. var (
  270. register = asm.Visit(ctx.RegisterR()).(int)
  271. instruction = bpf.LoadScratch{N: register}
  272. )
  273. switch target {
  274. case regA:
  275. instruction.Dst = bpf.RegA
  276. case regX:
  277. instruction.Dst = bpf.RegX
  278. }
  279. return instruction
  280. case ctx.IndirectX() != nil:
  281. var (
  282. number = asm.Visit(ctx.IndirectX()).(uint32)
  283. instruction = bpf.LoadIndirect{Off: number}
  284. )
  285. switch target {
  286. case regA, regX:
  287. instruction.Size = 1
  288. case reg2:
  289. instruction.Size = 2
  290. case reg4:
  291. instruction.Size = 4
  292. }
  293. return instruction
  294. case ctx.Extension() != nil:
  295. return bpf.LoadExtension{Num: asm.Visit(ctx.Extension()).(bpf.Extension)}
  296. default:
  297. return nil
  298. }
  299. }
  300. func (asm *assembler) VisitLoadOperator(ctx *parser.LoadOperatorContext) interface{} {
  301. switch {
  302. case ctx.LDA() != nil:
  303. return regA
  304. case ctx.LDX() != nil:
  305. return regX
  306. case ctx.LDB() != nil:
  307. return reg2
  308. case ctx.LDH() != nil:
  309. return reg4
  310. default:
  311. return invalid
  312. }
  313. }
  314. func (asm *assembler) VisitStoreOperation(ctx *parser.StoreOperationContext) interface{} {
  315. return bpf.StoreScratch{
  316. Src: asm.Visit(ctx.StoreOperator()).(bpf.Register),
  317. N: asm.Visit(ctx.RegisterR()).(int),
  318. }
  319. }
  320. func (asm *assembler) VisitStoreOperator(ctx *parser.StoreOperatorContext) interface{} {
  321. switch {
  322. case ctx.STA() != nil:
  323. return bpf.RegA
  324. case ctx.STX() != nil:
  325. return bpf.RegX
  326. default:
  327. return 0
  328. }
  329. }
  330. func (asm *assembler) VisitSimpleOperation(ctx *parser.SimpleOperationContext) interface{} {
  331. switch {
  332. case ctx.TAX() != nil:
  333. return bpf.TAX{}
  334. case ctx.TXA() != nil:
  335. return bpf.TXA{}
  336. case ctx.NEG() != nil:
  337. return bpf.NegateA{}
  338. default:
  339. return nil
  340. }
  341. }
  342. func (asm *assembler) VisitReturnOperation(ctx *parser.ReturnOperationContext) interface{} {
  343. switch {
  344. case ctx.RegisterA() != nil:
  345. return bpf.RetA{}
  346. case ctx.LiteralNumber() != nil:
  347. return bpf.RetConstant{Val: asm.Visit(ctx.LiteralNumber()).(uint32)}
  348. default:
  349. return nil
  350. }
  351. }
  352. func (asm *assembler) VisitNumber(ctx *parser.NumberContext) interface{} {
  353. return parseNumber(ctx.NUMBER().GetText())
  354. }
  355. func (asm *assembler) VisitLiteralNumber(ctx *parser.LiteralNumberContext) interface{} {
  356. return asm.Visit(ctx.Number())
  357. }
  358. func (asm *assembler) VisitIndirectX(ctx *parser.IndirectXContext) interface{} {
  359. return asm.Visit(ctx.Number())
  360. }
  361. func (asm *assembler) VisitAbsoluteNumber(ctx *parser.AbsoluteNumberContext) interface{} {
  362. return asm.Visit(ctx.Number())
  363. }
  364. func (asm *assembler) VisitRegisterA(ctx *parser.RegisterAContext) interface{} {
  365. return bpf.RegA
  366. }
  367. func (asm *assembler) VisitRegisterX(ctx *parser.RegisterXContext) interface{} {
  368. return bpf.RegX
  369. }
  370. func (asm *assembler) VisitRegisterR(ctx *parser.RegisterRContext) interface{} {
  371. return int(parseNumber(ctx.NUMBER().GetText()))
  372. }
  373. func (asm *assembler) VisitExtension(ctx *parser.ExtensionContext) interface{} {
  374. return parseExtension(ctx.IDENTIFIER().GetText())
  375. }