Class: Immunio::QueryTracker

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

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize ⇒ QueryTracker

Returns a new instance of QueryTracker.



367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
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
 
# File 'lib/immunio/plugins/active_record.rb', line 367

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.



447
448
449
450
451
452
453
454
455
456
 
# File 'lib/immunio/plugins/active_record.rb', line 447

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.



437
438
439
440
441
442
443
444
 
# File 'lib/immunio/plugins/active_record.rb', line 437

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.



568
569
570
571
572
573
574
575
576
577
578
579
 
# File 'lib/immunio/plugins/active_record.rb', line 568

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.



555
556
557
558
559
560
561
562
563
 
# File 'lib/immunio/plugins/active_record.rb', line 555

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.



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

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.



526
527
528
 
# File 'lib/immunio/plugins/active_record.rb', line 526

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

#call(payload) ⇒ Object

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



583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
 
# File 'lib/immunio/plugins/active_record.rb', line 583

def call(payload)
  Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
    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]
      params = relation_data[:params].clone
      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
      params = {}
      context_data = nil
      modifiers = {}
    end

    # Merge bound values
    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
        params[name] = value.to_s
      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,
                                                      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

#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.



503
504
505
 
# File 'lib/immunio/plugins/active_record.rb', line 503

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.



510
511
512
513
514
515
516
517
518
519
520
521
522
523
 
# File 'lib/immunio/plugins/active_record.rb', line 510

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



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

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



459
460
461
 
# File 'lib/immunio/plugins/active_record.rb', line 459

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

#reset(relation_id) ⇒ Object

Reset per-execution data for a relation.



646
647
648
649
650
651
652
 
# File 'lib/immunio/plugins/active_record.rb', line 646

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.



474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
 
# File 'lib/immunio/plugins/active_record.rb', line 474

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