Class: Immunio::QueryTracker

Inherits:
Object
  • Object
show all
Includes:
Singleton
Defined in:
lib/immunio/plugins/active_record.rb

Constant Summary collapse

DIALECTS =

Adapter name / db_dialect

{
  sqlite: 'sqlite3',
  postgres: 'postgres',
  mysql: 'mysql',
  mysql2: 'mysql',
  ibm_db: 'db2',
  oracle: 'oracle',
  oracleenhanced: 'oracle',
}.freeze

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize ⇒ QueryTracker

Returns a new instance of QueryTracker.



392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
 
# File 'lib/immunio/plugins/active_record.rb', line 392

def initialize
  # The data in these hashes represent relations and connections whose
  # lifecycle cannot be easily inferred. A relation could be kept around
  # across multiple HTTP requests, for example. We defined finalizers to
  # clean up data in the hashes when the objects they are linked to are
  # released by the Ruby runtime.
  #
  # Note: Relations have an associated connection and are not accessed by
  # other connections. Connections are associated with a thread and are not
  # accessed by other threads at the same time. Thus, there is no need for
  # thread safety in any of the logic in this class.

  # Data about a relation. The data inside is stored at different times and
  # must be reset properly:
  #
  # * params and relation_data: Added when ActiveRecord::Relation API calls
  #   are made, like #where. Should never be reset for a given relation.
  # * ast_data and modifiers: Added when a relation is converted
  #   into a SQL query statement. Should be reset after every query
  #   execution.
  @relation_data = Hash.new do |relation_data, relation_id|
    # This should never happen, but if it does it's a sign of an impending
    # memory leak. Log it, but just let it happen as it would be hard to
    # handle elsewhere if we did not set up the data.
    unless ObjectSpace._id2ref(relation_id).is_a? ActiveRecord::Relation
      name = if ObjectSpace._id2ref(relation_id).is_a? Class
        "#{ObjectSpace._id2ref(relation_id).name} Class"
      else
        "#{ObjectSpace._id2ref(relation_id).class.name} Instance"
      end
      Immunio.logger.warn {"Creating relation data for non-relation: #{name}"}
      Immunio.logger.debug {"Call stack:\n#{caller.join "\n"}"}
    end

    # NOTE: If you hold a reference to the relation here, like say:
    #
    # relation = ObjectSpace._id2ref(relation_id)
    #
    # the scope for the block will hold the relation and it will never be
    # released.
    ObjectSpace.define_finalizer ObjectSpace._id2ref(relation_id),
                                 self.class.finalize_relation(relation_id)

    relation_data[relation_id] = {
      params: {},
      relation_data: [],
      ast_data: [],
      modifiers: Hash.new do |modifiers, type|
        modifiers[type] = []
      end
    }
  end

  # Stacks of relations for each connection. Used to find the appropriate
  # relation for a connection when a query is executed. A stack is used
  # because some relation methods create new relations and call other
  # relation methods on the new relations.
  @relations = Hash.new do |relations, connection_id|
    connection = ObjectSpace._id2ref(connection_id)
    ObjectSpace.define_finalizer(connection, self.class.finalize_connection(connection_id))

    relations[connection_id] = []
  end

  # Last spawned relations for connections. Used for a hack to propagate
  # params to the right relation in Rails 3.
  @last_spawned_relations = {}
end

Class Method Details

.finalize_connection(connection_id) ⇒ Object

Delete a connection record when the connection object is released.



472
473
474
475
476
477
478
479
480
481
 
# File 'lib/immunio/plugins/active_record.rb', line 472

def self.finalize_connection(connection_id)
  proc do
    relations = instance.instance_variable_get(:@relations)

    # Check if key exists, delete will call the default value block if not
    relations.delete(connection_id) if relations.has_key? connection_id

    instance.instance_variable_get(:@last_spawned_relations).delete connection_id
  end
end

.finalize_relation(relation_id) ⇒ Object

Delete a relation record when the relation object is released.



462
463
464
465
466
467
468
469
 
# File 'lib/immunio/plugins/active_record.rb', line 462

def self.finalize_relation(relation_id)
  proc do
    relation_data = instance.instance_variable_get(:@relation_data)

    # Check if key exists, delete will call the default value block if not
    relation_data.delete(relation_id) if relation_data.has_key? relation_id
  end
end

Instance Method Details

#add_ast_data(ast_node_name, connection_id) ⇒ Object

Add data about an Arel AST node to the context data for the connection. This only occurs during conversion of a relation to SQL statement, so AST context data is never copied from one relation to another.



593
594
595
596
597
598
599
600
601
602
603
604
 
# File 'lib/immunio/plugins/active_record.rb', line 593

def add_ast_data(ast_node_name, connection_id)
  relation_id = @relations[connection_id].last

  # This can occur if the query statement was cached and there's no relation
  # associated with the connection. That's ok here, though, because there's
  # such a limited number of cacheable statement structures that we don't
  # need AST info to differentiate between queries with the same stack
  # trace.
  return unless relation_id

  @relation_data[relation_id][:ast_data] << "Arel AST visited node: #{ast_node_name}"
end

#add_modifier(type, value, connection_id) ⇒ Object

Add a modifier to the current relation for the connection. This only occurs during conversion of a relation to SQL statement, so modifiers are never copied from one relation to another.



580
581
582
583
584
585
586
587
588
 
# File 'lib/immunio/plugins/active_record.rb', line 580

def add_modifier(type, value, connection_id)
  relation_id = @relations[connection_id].last

  # This can occur if the query statement isn't generated by the app but by
  # ActiveRecord itself.
  return unless relation_id

  @relation_data[relation_id][:modifiers][type] << value
end

#add_param(name, value, connection_id) ⇒ Object

Add a parameter to the current relation for the connection. If the relation is copied or merged into another relation, the param will also be copied.



558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
 
# File 'lib/immunio/plugins/active_record.rb', line 558

def add_param(name, value, connection_id)
  relation_id = @relations[connection_id].last

  # This can occur if the query statement isn't generated by the app but by
  # ActiveRecord itself.
  return unless relation_id

  params = @relation_data[relation_id][:params]

  # If no name given, use index.
  if name.nil?
    name = params.size.to_s
  end

  Immunio.logger.debug { "Adding ActiveRecord SQL param to relation #{relation_id} (name: #{name}, value: #{value})" }

  params[name] = value
end

#add_relation_data(relation, data) ⇒ Object

Add relation API context data to the relation.



551
552
553
 
# File 'lib/immunio/plugins/active_record.rb', line 551

def add_relation_data(relation, data)
  @relation_data[relation.object_id][:relation_data] << data
end

#call(payload, adapter_name) ⇒ Object

Evaluate a SQL call. This occurs after Arel AST conversion of a relation to a statement.



623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
 
# File 'lib/immunio/plugins/active_record.rb', line 623

def call(payload, adapter_name)
  Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
    # The #{payload} in the log string is causing lots of Rails 3.2
    # upstream test failures with Ruby 1.9.3.
    # Immunio.logger.debug { "New ActiveRecord SQL query: #{payload}" }

    connection_id = payload[:connection_id]

    relation_id = @relations[connection_id].last

    if should_ignore? payload[:sql]
      Immunio.logger.debug { "Ignoring query as it was generated by ActiveRecord itself (#{payload[:sql]})" }
      return
    end

    if relation_id
      # Note: If a relation is released between when it is converted to a
      # SQL statement and now, we would lose the data and additionally leak
      # an empty entry in the @relation_data hash. I don't believe this is
      # possible due to how we wrap things, but there's no explicit
      # guarantee.
      relation_data = @relation_data[relation_id]
      context_data = (relation_data[:relation_data] + relation_data[:ast_data]).join "\n"

      # modifiers must be cloned because it will be cleared when the
      # relation is reset.
      modifiers = relation_data[:modifiers].clone
    else
      context_data = nil
      modifiers = {}
    end

    # Merge bound values into params
    params = {}
    question_marks = 0
    payload[:binds].each do |(column, value)|
      if column.nil?
        params["?:#{question_marks}"] = value.to_s
        question_marks = question_marks + 1
      else
        # When using the activerecord-sqlserver-adapter gem, the "column" is
        # the actual param name.
        name = column.respond_to?(:name) ? column.name : column.to_s

        # Some AR tests (sqlite3 in particular) initialize a QueryAttribute
        # with a nil `name`. This guards againt passing a nil key to Lua.
        params[name] = value.to_s if name
      end
    end

    strict_context, loose_context, stack = Immunio::Context.context context_data

    # Send in additional_context_data for debugging purposes
    Immunio.run_hook! "active_record",
                      "sql_execute",
                      sql: payload[:sql],
                      connection_uuid: connection_id.to_s,
                      db_dialect: db_dialect(adapter_name),
                      params: params,
                      modifiers: modifiers,
                      context_key: strict_context,
                      loose_context_key: loose_context,
                      stack: stack,
                      additional_context_data: context_data

    reset relation_id
  end
end

#db_dialect(adapter_name) ⇒ Object 



617
618
619
 
# File 'lib/immunio/plugins/active_record.rb', line 617

def db_dialect(adapter_name)
  DIALECTS.fetch(adapter_name.downcase.to_sym, 'unknown')
end

#last_spawned_relation(connection) ⇒ Object

Retrieve the last spawned relation for the connection. This is an ugly hack for a poor implementation of the #where and #having methods in Rails 3.



528
529
530
 
# File 'lib/immunio/plugins/active_record.rb', line 528

def last_spawned_relation(connection)
  ObjectSpace._id2ref @last_spawned_relations[connection.object_id]
end

#merge_relations(relation, other) ⇒ Object

Called when two relations are merged. The data for the other relation must be copied into the current relation. AST data and modifiers should be empty and don’t need to be copied.



535
536
537
538
539
540
541
542
543
544
545
546
547
548
 
# File 'lib/immunio/plugins/active_record.rb', line 535

def merge_relations(relation, other)
  params = @relation_data[relation.object_id][:params]
  other_params = @relation_data[other.object_id][:params]

  other_params.each_pair do |name, value|
    # Update numeric ID for current relation if name is an integer.
    name = params.size.to_s if name.to_i.to_s == name

    params[name] = value
  end

  other_data = @relation_data[other.object_id][:relation_data]
  @relation_data[relation.object_id][:relation_data] += other_data
end

#pop_relation(relation) ⇒ Object

Pop a relation off the stack for its connection



489
490
491
492
493
494
495
 
# File 'lib/immunio/plugins/active_record.rb', line 489

def pop_relation(relation)
  popped = @relations[relation.connection.object_id].pop
  unless popped == relation.object_id
    Immunio.logger.warn {"Popped wrong relation, expected: #{relation}, popped: #{popped}"}
    Immunio.logger.debug {"Call stack:\n#{caller.join "\n"}"}
  end
end

#push_relation(relation) ⇒ Object

Push a relation onto the stack for its connection



484
485
486
 
# File 'lib/immunio/plugins/active_record.rb', line 484

def push_relation(relation)
  @relations[relation.connection.object_id] << relation.object_id
end

#reset(relation_id) ⇒ Object

Reset per-execution data for a relation.



693
694
695
696
697
698
699
 
# File 'lib/immunio/plugins/active_record.rb', line 693

def reset(relation_id)
  return unless relation_id

  [:ast_data, :modifiers].each do |type|
    @relation_data[relation_id][type].clear
  end
end

#spawn_relation(orig, new) ⇒ Object

Called when a relation is cloned. The data for the new relation must also be copied from the old relation.



499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
 
# File 'lib/immunio/plugins/active_record.rb', line 499

def spawn_relation(orig, new)
  orig_id = orig.object_id
  new_id = new.object_id

  # If we weren't tracking the original relation, don't bother setting up
  # the new relation yet.
  if @relation_data.has_key? orig_id
    # ast_data and modifiers should be empty, but we must clone modifiers to
    # get the initializer.
    @relation_data[new_id] = {
      params: @relation_data[orig_id][:params].clone,
      relation_data: @relation_data[orig_id][:relation_data].clone,
      ast_data: @relation_data[orig_id][:ast_data] = [],
      modifiers: @relation_data[orig_id][:modifiers].clone
    }

    # The default block for the @relation_data hash isn't called when
    # assigning a value to a new key. We must set up the finalizer manually.
    ObjectSpace.define_finalizer(new, self.class.finalize_relation(new_id))
  end

  # Save the last spawned relation for a hack for storing params from #where
  # and #having.
  @last_spawned_relations[new.connection.object_id] = new_id
end